Upload conformer_model.py with huggingface_hub

conformer_model.py

@@ -0,0 +1,626 @@
+"""The conformer that turns prosody features into a mel spectrogram.
+
+Extracted verbatim from the training code -- `prosody_embeddings_to_mel` --
+so the architecture here is the one the checkpoint was trained with rather
+than a reimplementation of it. Only what inference touches: the model
+classes and `generate()`, with no Trainer, dataset or metric code.
+
+`generate()` is non-autoregressive. It sizes its output from the real input
+length and takes a padding mask, so a phrase is predicted at its own
+duration -- which is why this is the path the published numbers come from.
+The ONNX export cannot do that: its input width is baked in at 2000 frames,
+so every phrase is padded to ten seconds and the decoder attends over the
+padding. Measured against the training pipeline's own predictions, 0.9985
+correlation here against 0.9357 through the export.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint as grad_checkpoint
+
+
+class SinusoidalPositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x):
+        # x is [B, T, D] (batch_first)
+        return x + self.pe[: x.size(1)]
+
+
+# ----------------------------
+# Conformer Block
+# ----------------------------
+
+
+class ConformerBlock(nn.Module):
+    """Single conformer block: FFN -> Self-Attention -> Conv -> FFN.
+
+    Combines local convolutions with self-attention for both local
+    detail and global context, at lower cost than a full transformer
+    decoder (conv handles local, attention handles global).
+    """
+    def __init__(self, d_model, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1):
+        super().__init__()
+        dim_ff = dim_feedforward or d_model * 4
+
+        # First half-step FFN
+        self.ffn1_norm = nn.LayerNorm(d_model)
+        self.ffn1 = nn.Sequential(
+            nn.Linear(d_model, dim_ff), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim_ff, d_model), nn.Dropout(dropout),
+        )
+
+        # Self-attention
+        self.attn_norm = nn.LayerNorm(d_model)
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=True)
+        self.attn_dropout = nn.Dropout(dropout)
+
+        # Depthwise convolution module
+        self.conv_norm = nn.LayerNorm(d_model)
+        self.conv = nn.Sequential(
+            nn.Conv1d(d_model, d_model, 1),  # pointwise
+            nn.GELU(),
+            nn.Conv1d(d_model, d_model, kernel_size, padding=kernel_size // 2, groups=d_model),  # depthwise
+            nn.LayerNorm(d_model),  # NOTE: applied after transpose in forward
+            nn.GELU(),
+            nn.Conv1d(d_model, d_model, 1),  # pointwise
+            nn.Dropout(dropout),
+        )
+        # Replace the LayerNorm inside Sequential with identity — we apply it manually
+        self.conv_ln = self.conv[3]
+        self.conv[3] = nn.Identity()
+
+        # Second half-step FFN
+        self.ffn2_norm = nn.LayerNorm(d_model)
+        self.ffn2 = nn.Sequential(
+            nn.Linear(d_model, dim_ff), nn.GELU(), nn.Dropout(dropout),
+            nn.Linear(dim_ff, d_model), nn.Dropout(dropout),
+        )
+
+        self.final_norm = nn.LayerNorm(d_model)
+
+    def forward(self, x, padding_mask=None):
+        """
+        Args:
+            x: [B, T, D]
+            padding_mask: [B, T] with True for padding
+        Returns:
+            [B, T, D]
+        """
+        # Half-step FFN
+        x = x + 0.5 * self.ffn1(self.ffn1_norm(x))
+
+        # Self-attention
+        h = self.attn_norm(x)
+        attn_out, _ = self.self_attn(h, h, h, key_padding_mask=padding_mask)
+        x = x + self.attn_dropout(attn_out)
+
+        # Conv module
+        h = self.conv_norm(x)
+        h = h.transpose(1, 2)  # [B, D, T]
+        # Apply conv layers, but handle LayerNorm which needs [B, T, D]
+        h = self.conv[0](h)  # pointwise
+        h = self.conv[1](h)  # GELU
+        h = self.conv[2](h)  # depthwise
+        h = h.transpose(1, 2)  # [B, T, D]
+        h = self.conv_ln(h)   # LayerNorm
+        h = self.conv[4](h.transpose(1, 2))  # GELU (expects [B, D, T] from Conv1d but GELU is elementwise)
+        h = self.conv[5](h)  # pointwise
+        h = self.conv[6](h)  # dropout
+        h = h.transpose(1, 2)  # [B, T, D]
+        x = x + h
+
+        # Half-step FFN
+        x = x + 0.5 * self.ffn2(self.ffn2_norm(x))
+
+        x = self.final_norm(x)
+
+        if padding_mask is not None:
+            x = x.masked_fill(padding_mask.unsqueeze(-1), 0.0)
+
+        return x
+
+
+class ConformerDecoder(nn.Module):
+    """Stack of conformer blocks as a decoder."""
+    def __init__(self, d_model, num_layers=4, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1):
+        super().__init__()
+        self.blocks = nn.ModuleList([
+            ConformerBlock(d_model, nhead, kernel_size, dim_feedforward, dropout)
+            for _ in range(num_layers)
+        ])
+
+    def forward(self, x, padding_mask=None):
+        for block in self.blocks:
+            x = block(x, padding_mask)
+        return x
+
+
+# ----------------------------
+# Model
+# ----------------------------
+
+
+class WindowedConformerBlock(ConformerBlock):
+    """ConformerBlock with windowed self-attention for O(T*W) memory instead of O(T²).
+
+    Splits the sequence into non-overlapping windows of size W and applies
+    self-attention independently within each window. The convolution module
+    still provides cross-window local context.
+    """
+    def __init__(self, d_model, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1, attn_window=256):
+        super().__init__(d_model, nhead, kernel_size, dim_feedforward, dropout)
+        self.attn_window = attn_window
+
+    def forward(self, x, padding_mask=None):
+        """Same as ConformerBlock but with windowed attention."""
+        # Half-step FFN
+        x = x + 0.5 * self.ffn1(self.ffn1_norm(x))
+
+        # Windowed self-attention
+        h = self.attn_norm(x)
+        attn_out = self._windowed_attention(h, padding_mask)
+        x = x + self.attn_dropout(attn_out)
+
+        # Conv module (unchanged — provides cross-window context)
+        h = self.conv_norm(x)
+        h = h.transpose(1, 2)  # [B, D, T]
+        h = self.conv[0](h)  # pointwise
+        h = self.conv[1](h)  # GELU
+        h = self.conv[2](h)  # depthwise
+        h = h.transpose(1, 2)  # [B, T, D]
+        h = self.conv_ln(h)   # LayerNorm
+        h = self.conv[4](h.transpose(1, 2))  # GELU
+        h = self.conv[5](h)  # pointwise
+        h = self.conv[6](h)  # dropout
+        h = h.transpose(1, 2)  # [B, T, D]
+        x = x + h
+
+        # Half-step FFN
+        x = x + 0.5 * self.ffn2(self.ffn2_norm(x))
+
+        x = self.final_norm(x)
+
+        if padding_mask is not None:
+            x = x.masked_fill(padding_mask.unsqueeze(-1), 0.0)
+
+        return x
+
+    def _windowed_attention(self, x, padding_mask):
+        """Self-attention within non-overlapping windows.
+
+        Args:
+            x: [B, T, D] — already normed
+            padding_mask: [B, T] with True for padding, or None
+
+        Returns:
+            [B, T, D]
+        """
+        B, T, D = x.shape
+        W = self.attn_window
+
+        # Short sequences: full attention (no windowing needed)
+        if T <= W:
+            out, _ = self.self_attn(x, x, x, key_padding_mask=padding_mask)
+            return out
+
+        # Pad sequence to a multiple of W
+        pad_len = (W - T % W) % W
+        if pad_len > 0:
+            x = torch.nn.functional.pad(x, (0, 0, 0, pad_len))  # pad time dim
+            if padding_mask is not None:
+                padding_mask = torch.nn.functional.pad(padding_mask, (0, pad_len), value=True)
+
+        T_padded = x.size(1)
+        num_windows = T_padded // W
+
+        # Reshape: [B, num_windows, W, D] -> [B*num_windows, W, D]
+        x_win = x.reshape(B * num_windows, W, D)
+
+        # Reshape mask: [B, num_windows, W] -> [B*num_windows, W]
+        mask_win = None
+        if padding_mask is not None:
+            mask_win = padding_mask.reshape(B * num_windows, W)
+            # If an entire window is padding, MHA will produce NaN.
+            # Detect all-padding windows and handle them.
+            all_padding = mask_win.all(dim=1)  # [B*num_windows]
+            if all_padding.any():
+                # Temporarily unmask one position so MHA doesn't get empty input
+                mask_win = mask_win.clone()
+                mask_win[all_padding, 0] = False
+
+        attn_out, _ = self.self_attn(x_win, x_win, x_win, key_padding_mask=mask_win)
+
+        # Zero out windows that were entirely padding
+        if padding_mask is not None and all_padding.any():
+            attn_out[all_padding] = 0.0
+
+        # Reshape back: [B*num_windows, W, D] -> [B, T_padded, D] -> [B, T, D]
+        attn_out = attn_out.reshape(B, T_padded, D)[:, :T, :]
+        return attn_out
+
+
+class WindowedConformerDecoder(nn.Module):
+    """Stack of windowed conformer blocks."""
+    def __init__(self, d_model, num_layers=4, nhead=4, kernel_size=5, dim_feedforward=None, dropout=0.1, attn_window=256):
+        super().__init__()
+        self.blocks = nn.ModuleList([
+            WindowedConformerBlock(d_model, nhead, kernel_size, dim_feedforward, dropout, attn_window)
+            for _ in range(num_layers)
+        ])
+
+    def forward(self, x, padding_mask=None):
+        for block in self.blocks:
+            x = block(x, padding_mask)
+        return x
+
+
+# ----------------------------
+# Convolutional Encoder
+# ----------------------------
+
+
+class ConvEncoder(nn.Module):
+    """Multi-layer 1D convolutional encoder with stride-2 downsampling.
+
+    Designed for smooth, low-dimensional prosodic signals (f0 + intensity).
+    Uses dilated convolutions to capture multi-scale temporal context efficiently,
+    with a final stride-2 layer to handle the 200 Hz → 100 Hz rate conversion.
+    """
+    def __init__(self, input_dim, d_model, num_layers=4, kernel_size=7, dropout=0.1):
+        super().__init__()
+        self.input_proj = nn.Conv1d(input_dim, d_model, kernel_size=1)
+
+        layers = []
+        for i in range(num_layers - 1):
+            dilation = 2 ** i  # 1, 2, 4, 8... expanding receptive field
+            padding = (kernel_size - 1) * dilation // 2
+            layers.append(ConvBlock(d_model, kernel_size, dilation, padding, dropout))
+
+        self.layers = nn.ModuleList(layers)
+
+        # Final stride-2 conv for downsampling (200 Hz → 100 Hz)
+        self.downsample = nn.Conv1d(
+            d_model, d_model,
+            kernel_size=4, stride=2, padding=1,
+        )
+        self.downsample_norm = nn.LayerNorm(d_model)
+
+    def forward(self, x, padding_mask=None):
+        """
+        Args:
+            x: [B, T, input_dim]
+            padding_mask: [B, T] with True for padding positions
+
+        Returns:
+            out: [B, T//2, d_model]
+            out_mask: [B, T//2] or None
+        """
+        # [B, T, C] → [B, C, T] for Conv1d
+        h = x.transpose(1, 2)
+        h = self.input_proj(h)
+
+        # Zero out padded positions before conv
+        if padding_mask is not None:
+            h = h.masked_fill(padding_mask.unsqueeze(1), 0.0)
+
+        for layer in self.layers:
+            h = layer(h, padding_mask)
+
+        # Stride-2 downsampling
+        h = self.downsample(h)  # [B, d_model, T//2]
+        h = h.transpose(1, 2)  # [B, T//2, d_model]
+        h = self.downsample_norm(h)
+
+        # Compute output padding mask
+        out_mask = None
+        if padding_mask is not None:
+            T_out = h.size(1)
+            # A downsampled position is padding if both source frames were padding
+            T_in = padding_mask.size(1)
+            if T_in % 2 != 0:
+                padding_mask = torch.cat([
+                    padding_mask,
+                    torch.ones(padding_mask.size(0), 1, dtype=torch.bool, device=padding_mask.device)
+                ], dim=1)
+            mask_pairs = padding_mask.view(padding_mask.size(0), -1, 2)
+            out_mask = mask_pairs.all(dim=-1)[:, :T_out]
+
+        if out_mask is not None:
+            h = h.masked_fill(out_mask.unsqueeze(-1), 0.0)
+
+        return h, out_mask
+
+
+class ConvBlock(nn.Module):
+    """Residual dilated conv block with pre-norm."""
+    def __init__(self, d_model, kernel_size, dilation, padding, dropout):
+        super().__init__()
+        self.norm = nn.LayerNorm(d_model)
+        self.conv = nn.Conv1d(
+            d_model, d_model,
+            kernel_size=kernel_size,
+            dilation=dilation,
+            padding=padding,
+        )
+        self.activation = nn.GELU()
+        self.dropout = nn.Dropout(dropout)
+        self.proj = nn.Conv1d(d_model, d_model, kernel_size=1)
+
+    def forward(self, x, padding_mask=None):
+        """
+        Args:
+            x: [B, d_model, T]
+            padding_mask: [B, T] with True for padding
+        Returns:
+            [B, d_model, T]
+        """
+        residual = x
+        # Pre-norm (need to transpose for LayerNorm)
+        h = x.transpose(1, 2)  # [B, T, D]
+        h = self.norm(h)
+        h = h.transpose(1, 2)  # [B, D, T]
+
+        h = self.conv(h)
+        h = self.activation(h)
+        h = self.dropout(h)
+        h = self.proj(h)
+        h = self.dropout(h)
+
+        out = residual + h
+
+        if padding_mask is not None:
+            out = out.masked_fill(padding_mask.unsqueeze(1), 0.0)
+
+        return out
+
+
+# ----------------------------
+# Model
+# ----------------------------
+
+
+class FeaturesToMelConv(nn.Module):
+    """Non-autoregressive conv-encoder + Conformer-decoder for f0+intensity to mel.
+
+    Architecture rationale:
+    - Conv encoder: F0 and intensity are smooth, slowly-varying 1D signals.
+      Dilated convolutions naturally capture multi-scale temporal patterns
+      (pitch contours, intensity envelopes) with the right inductive bias.
+      The stride-2 final layer handles the 200→100 Hz rate conversion
+      as part of the encoding, rather than as a separate step.
+    - Conformer decoder: Refines the downsampled features using both local
+      (depthwise conv) and global (self-attention) context to produce
+      the mel spectrogram. Self-attention here is justified because the
+      decoder operates at 100 Hz on d_model-dimensional features — a
+      richer representation where global dependencies matter.
+
+    Memory optimization: gradient checkpointing on decoder blocks (recomputes
+    activations during backward instead of storing them).
+    """
+    def __init__(
+        self,
+        input_dim=2,
+        mel_dim=14,
+        d_model=256,
+        nhead=4,
+        num_encoder_layers=4,
+        num_decoder_layers=4,
+        encoder_kernel_size=7,
+        conv_kernel_size=5,
+        dim_feedforward=None,
+        dropout=0.1,
+        max_len=16000,
+        loss_type="mse",
+        length_ratio=0.5,
+        gradient_checkpointing=True,
+        decoder_only=False,
+        attn_window=0,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.mel_dim = mel_dim
+        self.max_len = max_len
+        self.length_ratio = length_ratio
+        self.gradient_checkpointing = gradient_checkpointing
+        self.decoder_only = decoder_only
+
+        if not decoder_only:
+            self.encoder = ConvEncoder(
+                input_dim=input_dim,
+                d_model=d_model,
+                num_layers=num_encoder_layers,
+                kernel_size=encoder_kernel_size,
+                dropout=dropout,
+            )
+        else:
+            # Decoder-only: simple linear projection + learned stride-2 downsample
+            self.input_proj = nn.Linear(input_dim, d_model)
+            self.downsample = nn.Linear(d_model * 2, d_model)
+
+        self.pos_decoder = SinusoidalPositionalEncoding(d_model, max_len)
+
+        dim_ff = dim_feedforward if dim_feedforward is not None else d_model * 4
+        if attn_window > 0:
+            self.decoder = WindowedConformerDecoder(
+                d_model=d_model,
+                num_layers=num_decoder_layers,
+                nhead=nhead,
+                kernel_size=conv_kernel_size,
+                dim_feedforward=dim_ff,
+                dropout=dropout,
+                attn_window=attn_window,
+            )
+        else:
+            self.decoder = ConformerDecoder(
+                d_model=d_model,
+                num_layers=num_decoder_layers,
+                nhead=nhead,
+                kernel_size=conv_kernel_size,
+                dim_feedforward=dim_ff,
+                dropout=dropout,
+            )
+
+        self.output_proj = nn.Linear(d_model, mel_dim)
+
+        if loss_type == "mse":
+            self.criterion = nn.MSELoss(reduction="none")
+        elif loss_type == "l1":
+            self.criterion = nn.L1Loss(reduction="none")
+        elif loss_type == "smoothl1":
+            self.criterion = nn.SmoothL1Loss(reduction="none")
+        else:
+            raise ValueError(f"Unknown loss_type={loss_type}")
+
+    def _run_decoder(self, tgt, padding_mask):
+        """Run decoder with optional gradient checkpointing per block."""
+        if self.gradient_checkpointing and self.training:
+            for block in self.decoder.blocks:
+                tgt = grad_checkpoint(block, tgt, padding_mask, use_reentrant=False)
+            return tgt
+        else:
+            return self.decoder(tgt, padding_mask=padding_mask)
+
+    def _encode(self, inputs, input_padding_mask):
+        """Encode inputs, returns (encoded, mask) at output frame rate."""
+        if not self.decoder_only:
+            return self.encoder(inputs, input_padding_mask)
+
+        # Decoder-only path: project + pair-merge downsample
+        B, T, _ = inputs.shape
+        h = self.input_proj(inputs)  # [B, T, d_model]
+
+        # Pad to even length
+        if T % 2 != 0:
+            h = torch.cat([h, torch.zeros(B, 1, self.d_model, device=h.device)], dim=1)
+            if input_padding_mask is not None:
+                input_padding_mask = torch.cat([
+                    input_padding_mask,
+                    torch.ones(B, 1, dtype=torch.bool, device=h.device)
+                ], dim=1)
+            T = T + 1
+
+        # Merge pairs: [B, T, D] -> [B, T//2, 2*D] -> [B, T//2, D]
+        h = h.view(B, T // 2, 2 * self.d_model)
+        h = self.downsample(h)
+
+        # Compute output mask
+        out_mask = None
+        if input_padding_mask is not None:
+            mask_pairs = input_padding_mask.view(B, T // 2, 2)
+            out_mask = mask_pairs.all(dim=-1)
+
+        if out_mask is not None:
+            h = h.masked_fill(out_mask.unsqueeze(-1), 0.0)
+
+        return h, out_mask
+
+    def forward(self, inputs, labels, input_padding_mask=None, label_padding_mask=None):
+        """
+        Args:
+            inputs: [B, T_src, input_dim] - f0 + intensity at 200 Hz
+            labels: [B, T_tgt, mel_dim] - ground truth mel at 100 Hz
+            input_padding_mask: [B, T_src] with True for padding
+            label_padding_mask: [B, T_tgt] with True for padding
+
+        Returns:
+            dict with 'loss' and 'logits'
+        """
+        B = labels.size(0)
+        T_tgt = labels.size(1)
+
+        # Encode + downsample
+        enc_out, enc_mask = self._encode(inputs, input_padding_mask)
+
+        # Align to target length
+        if enc_out.size(1) < T_tgt:
+            pad_len = T_tgt - enc_out.size(1)
+            enc_out = torch.cat([enc_out, torch.zeros(B, pad_len, self.d_model, device=enc_out.device)], dim=1)
+            if enc_mask is not None:
+                enc_mask = torch.cat([enc_mask, torch.ones(B, pad_len, dtype=torch.bool, device=enc_out.device)], dim=1)
+        elif enc_out.size(1) > T_tgt:
+            enc_out = enc_out[:, :T_tgt, :]
+            if enc_mask is not None:
+                enc_mask = enc_mask[:, :T_tgt]
+
+        # Add positional encoding for decoder
+        tgt = enc_out * math.sqrt(self.d_model)
+        tgt = self.pos_decoder(tgt)
+
+        # Conformer decoder refines the features
+        output = self._run_decoder(tgt, label_padding_mask)
+
+        logits = self.output_proj(output)  # [B, T_tgt, mel_dim]
+
+        # Compute loss
+        loss = self.criterion(logits, labels)
+        if label_padding_mask is not None:
+            active_mask = ~label_padding_mask
+            denom = (active_mask.sum() * logits.size(-1)).clamp_min(1).to(loss.dtype)
+            loss = (loss * active_mask.unsqueeze(-1)).sum() / (denom + 1e-8)
+        else:
+            loss = loss.mean()
+
+        return {"loss": loss, "logits": logits}
+
+    def generate(self, inputs, input_padding_mask=None, max_steps=None, length_ratio=None, return_lengths=False):
+        """Non-autoregressive parallel generation."""
+        B = inputs.size(0)
+        device = inputs.device
+
+        if length_ratio is None:
+            length_ratio = self.length_ratio
+
+        if max_steps is None:
+            if input_padding_mask is not None:
+                input_lengths = (~input_padding_mask).sum(dim=1)
+                target_lengths = (input_lengths.float() * length_ratio).round().long()
+                T_gen = int(target_lengths.max().item())
+            else:
+                input_lengths = torch.full((B,), inputs.size(1), device=device, dtype=torch.long)
+                target_lengths = (input_lengths.float() * length_ratio).round().long()
+                T_gen = int(target_lengths[0].item())
+        else:
+            T_gen = max_steps
+            target_lengths = torch.full((B,), T_gen, device=device, dtype=torch.long)
+
+        T_gen = max(1, min(T_gen, self.max_len))
+
+        output_padding_mask = torch.zeros(B, T_gen, dtype=torch.bool, device=device)
+        for i in range(B):
+            if target_lengths[i] < T_gen:
+                output_padding_mask[i, target_lengths[i]:] = True
+
+        enc_out, _ = self._encode(inputs, input_padding_mask)
+
+        if enc_out.size(1) < T_gen:
+            pad_len = T_gen - enc_out.size(1)
+            enc_out = torch.cat([enc_out, torch.zeros(B, pad_len, self.d_model, device=device)], dim=1)
+        elif enc_out.size(1) > T_gen:
+            enc_out = enc_out[:, :T_gen, :]
+
+        tgt = enc_out * math.sqrt(self.d_model)
+        tgt = self.pos_decoder(tgt)
+
+        output = self.decoder(tgt, padding_mask=output_padding_mask)
+        generated = self.output_proj(output)
+
+        if return_lengths:
+            return generated, target_lengths, output_padding_mask
+        return generated
+
+
+# ----------------------------
+# Dataset
+# ----------------------------