prosody2text-weights / conformer_model.py

Upload conformer_model.py with huggingface_hub

12d4f2a verified 6 days ago

22.8 kB

"""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
# ----------------------------

	"""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
	# ----------------------------