roformer.py

from typing import Tuple, Optional, List
import mlx.core as mx
import mlx.nn as nn
from .modules import RMSNorm, RotaryEmbedding, Transformer, MLP, ConformerConvModule, MacaronFF, Attention, Conformer
from .spectral import STFT, iSTFT

class BandSplit(nn.Module):
    def __init__(self, dim, dim_inputs: Tuple[int, ...]):
        super().__init__()
        # Matches hierarchy: to_features.layers.N.layers.0 (RMSNorm), .1 (Linear)
        self.to_features = nn.Sequential(*[
            nn.Sequential(RMSNorm(dim_in), nn.Linear(dim_in, dim)) 
            for dim_in in dim_inputs
        ])
        curr = 0; self.offsets = [0]
        for d in dim_inputs: curr += d; self.offsets.append(int(curr))
    def __call__(self, x):
        return mx.stack([layer(x[:, :, self.offsets[i]:self.offsets[i+1]]) for i, layer in enumerate(self.to_features.layers)], axis=2)

class MaskEstimator(nn.Module):
    def __init__(self, dim, num_channels, freqs_per_bands, depth=2):
        super().__init__()
        self.to_freqs = []
        for num_freqs in freqs_per_bands:
            out_dim = num_freqs * num_channels * 2
            # Matches hierarchy: to_freqs.layers.N.layers.0 (Sequential MLP), .1 (GLU)
            layers = MLP(dim, out_dim * 2, dim_hidden=dim*4, depth=depth)
            mlp_seq = nn.Sequential(*layers)
            self.to_freqs.append(nn.Sequential(mlp_seq, nn.GLU(axis=-1)))
        self.to_freqs = nn.Sequential(*self.to_freqs)
    def __call__(self, x):
        outputs = []
        for i, layer in enumerate(self.to_freqs.layers):
            band_input = x[:, :, i, :]
            outputs.append(layer(band_input))
        return mx.concatenate(outputs, axis=-1)

class BSRoformer(nn.Module):
    def __init__(self, dim, depth, stereo=True, num_stems=1, time_transformer_depth=1, freq_transformer_depth=1, linear_transformer_depth=0, heads=8, dim_head=64, mult=4, dropout=0.0, freqs_per_bands: Tuple[int, ...] = (2,) * 60, stft_n_fft=2048, stft_hop_length=441, stft_win_length=2048, zero_dc=True, mask_estimator_depth=2, stft_normalized=False, **kwargs):
        super().__init__()
        self.audio_channels = 2 if stereo else 1
        self.num_stems = num_stems
        dim_in = tuple(2 * f * self.audio_channels for f in freqs_per_bands)
        self.band_split, rope = BandSplit(dim, dim_in), RotaryEmbedding(dim_head)
        self.layers = nn.Sequential(*[nn.Sequential(Transformer(dim, time_transformer_depth, heads, dim_head, mult, dropout, rope), Transformer(dim, freq_transformer_depth, heads, dim_head, mult, dropout, rope)) for _ in range(depth)])
        self.final_norm = RMSNorm(dim)
        self.mask_estimators = nn.Sequential(*[
            MaskEstimator(dim, self.audio_channels, freqs_per_bands, depth=mask_estimator_depth)
            for _ in range(num_stems)
        ])
        self.zero_dc = zero_dc
        self.stft_normalized = stft_normalized
        
        # Native STFT/iSTFT
        self.stft = STFT(stft_n_fft, stft_hop_length, stft_win_length)
        self.istft = iSTFT(stft_n_fft, stft_hop_length, stft_win_length)
        
        # Compile the backbone
        self.fast_compute_masks = mx.compile(self._compute_masks)

    def _compute_masks(self, x):
        # Band split
        x = self.band_split(x)  # (B, Tf, num_bands, D)
        
        # Transformer blocks
        for lp in self.layers.layers:
            tt, ft = lp.layers
            Bc, Tc, Fc, Dc = x.shape
            x = tt(x.transpose(0, 2, 1, 3).reshape(Bc * Fc, Tc, Dc)).reshape(Bc, Fc, Tc, Dc).transpose(0, 2, 1, 3)
            x = ft(x.reshape(Bc * Tc, Fc, Dc)).reshape(Bc, Tc, Fc, Dc)
        x = self.final_norm(x)
        
        # Output masks
        masks = []
        for estimator in self.mask_estimators.layers:
            masks.append(estimator(x))
            
        return mx.stack(masks, axis=1) # (B, NumStems, Tf, F*S*2)

    def __call__(self, raw_audio):
        if raw_audio.ndim == 2: raw_audio = raw_audio[:, None, :]  # (B, 1, T) or (1, C, T)?
        if raw_audio.ndim == 2: raw_audio = raw_audio[None, :, :]
             
        B, S, T = raw_audio.shape
        x = raw_audio.reshape(B * S, T)
        stft_repr = self.stft(x)
        
        if self.stft_normalized:
            stft_repr = stft_repr * (self.stft.n_fft ** -0.5)

        if self.zero_dc:
            F = stft_repr.shape[1]
            mask_list = [mx.zeros((1, 1, 1)), mx.ones((F - 1, 1, 1))]
            mask = mx.concatenate(mask_list, axis=0) # (F, 1, 1)
            stft_repr = stft_repr * mask
            
        F = stft_repr.shape[1]
        Tf = stft_repr.shape[2]
        
        stft_repr = stft_repr.reshape(B, S, F, Tf, 2)
        x_backbone = stft_repr.transpose(0, 3, 2, 1, 4).reshape(B, Tf, F * S * 2)
        
        masks = self.fast_compute_masks(x_backbone)
        masks = masks.reshape(B, self.num_stems, Tf, F, S, 2)
        
        stft_expanded = stft_repr.transpose(0, 3, 2, 1, 4)[:, None, ...]
        
        spec_real, spec_imag = stft_expanded[..., 0], stft_expanded[..., 1]
        mask_real, mask_imag = masks[..., 0], masks[..., 1]
        
        out_real = spec_real * mask_real - spec_imag * mask_imag
        out_imag = spec_real * mask_imag + spec_imag * mask_real
        
        stft_masked = mx.stack([out_real, out_imag], axis=-1)
        stft_masked = stft_masked.transpose(0, 1, 4, 3, 2, 5)
        stft_masked_flat = stft_masked.reshape(B * self.num_stems * S, F, Tf, 2)
        
        if self.zero_dc:
             F_stft = stft_masked_flat.shape[1]
             mask_list = [mx.zeros((1, 1, 1)), mx.ones((F_stft - 1, 1, 1))]
             mask = mx.concatenate(mask_list, axis=0)
             stft_masked_flat = stft_masked_flat * mask

        if self.stft_normalized:
            stft_masked_flat = stft_masked_flat * (self.stft.n_fft ** 0.5)

        recon = self.istft(stft_masked_flat, length=T) 
        return recon.reshape(B, self.num_stems, S, T)

class BSConformer(nn.Module):
    def __init__(self, dim, depth, stereo=True, num_stems=1, time_conformer_depth=1, freq_conformer_depth=1, linear_conformer_depth=0, heads=8, dim_head=64, ff_mult=4, dropout=0.0, freqs_per_bands: Tuple[int, ...] = (2,) * 60, stft_n_fft=2048, stft_hop_length=441, stft_win_length=2048, zero_dc=True, mask_estimator_depth=2, conv_expansion_factor=2, conv_kernel_size=31, stft_normalized=False, **kwargs):
        super().__init__()
        self.audio_channels = 2 if stereo else 1
        self.num_stems = num_stems
        dim_in = tuple(2 * f * self.audio_channels for f in freqs_per_bands)
        self.band_split, rope = BandSplit(dim, dim_in), RotaryEmbedding(dim_head)
        
        blocks = []
        for _ in range(depth):
            block = []
            if linear_conformer_depth > 0:
                block.append(Transformer(dim, linear_conformer_depth, heads, dim_head, ff_mult, dropout, rope))
            
            block.append(Conformer(dim, time_conformer_depth, heads, dim_head, ff_mult, dropout, dropout, conv_expansion_factor, conv_kernel_size, rope))
            block.append(Conformer(dim, freq_conformer_depth, heads, dim_head, ff_mult, dropout, dropout, conv_expansion_factor, conv_kernel_size, rope))
            blocks.append(nn.Sequential(*block))
            
        self.layers = nn.Sequential(*blocks)
        self.final_norm = RMSNorm(dim)
        self.mask_estimators = nn.Sequential(*[
            MaskEstimator(dim, self.audio_channels, freqs_per_bands, depth=mask_estimator_depth)
            for _ in range(num_stems)
        ])
        self.zero_dc = zero_dc
        self.stft_normalized = stft_normalized
        
        self.stft = STFT(stft_n_fft, stft_hop_length, stft_win_length)
        self.istft = iSTFT(stft_n_fft, stft_hop_length, stft_win_length)
        
        self.fast_compute_masks = mx.compile(self._compute_masks)

    def _compute_masks(self, x):
        x = self.band_split(x)
        for lp in self.layers.layers:
            sublayers = lp.layers
            if len(sublayers) == 3:
                linear, tt, ft = sublayers
                Bc, Tc, Fc, Dc = x.shape
                x = linear(x.transpose(0, 2, 1, 3).reshape(Bc * Fc, Tc, Dc)).reshape(Bc, Fc, Tc, Dc).transpose(0, 2, 1, 3)
                x = tt(x.transpose(0, 2, 1, 3).reshape(Bc * Fc, Tc, Dc)).reshape(Bc, Fc, Tc, Dc).transpose(0, 2, 1, 3)
                x = ft(x.reshape(Bc * Tc, Fc, Dc)).reshape(Bc, Tc, Fc, Dc)
            else:
                tt, ft = sublayers
                Bc, Tc, Fc, Dc = x.shape
                x = tt(x.transpose(0, 2, 1, 3).reshape(Bc * Fc, Tc, Dc)).reshape(Bc, Fc, Tc, Dc).transpose(0, 2, 1, 3)
                x = ft(x.reshape(Bc * Tc, Fc, Dc)).reshape(Bc, Tc, Fc, Dc)
        
        x = self.final_norm(x)
        masks = [estimator(x) for estimator in self.mask_estimators.layers]
        return mx.stack(masks, axis=1)

    def __call__(self, raw_audio):
        if raw_audio.ndim == 2: raw_audio = raw_audio[None, :, :]
        B, S, T = raw_audio.shape
        x = raw_audio.reshape(B * S, T)
        stft_repr = self.stft(x)
        
        if self.stft_normalized:
             stft_repr = stft_repr * (self.stft.n_fft ** -0.5)
        
        if self.zero_dc:
            F = stft_repr.shape[1]
            mask_list = [mx.zeros((1, 1, 1)), mx.ones((F - 1, 1, 1))]
            mask = mx.concatenate(mask_list, axis=0)
            stft_repr = stft_repr * mask

        F, Tf = stft_repr.shape[1], stft_repr.shape[2]
        stft_repr = stft_repr.reshape(B, S, F, Tf, 2)
        x_backbone = stft_repr.transpose(0, 3, 2, 1, 4).reshape(B, Tf, F * S * 2)
        
        masks = self.fast_compute_masks(x_backbone)
        masks = masks.reshape(B, self.num_stems, Tf, F, S, 2)
        
        stft_expanded = stft_repr.transpose(0, 3, 2, 1, 4)[:, None, ...]
        
        spec_real, spec_imag = stft_expanded[..., 0], stft_expanded[..., 1]
        mask_real, mask_imag = masks[..., 0], masks[..., 1]
        
        out_real = spec_real * mask_real - spec_imag * mask_imag
        out_imag = spec_real * mask_imag + spec_imag * mask_real
        
        stft_masked = mx.stack([out_real, out_imag], axis=-1).transpose(0, 1, 4, 3, 2, 5)
        stft_masked_flat = stft_masked.reshape(B * self.num_stems * S, F, Tf, 2)
        
        if self.stft_normalized:
             stft_masked_flat = stft_masked_flat * (self.stft.n_fft ** 0.5)
        
        recon = self.istft(stft_masked_flat, length=T)
        return recon.reshape(B, self.num_stems, S, T)