mel_band_roformer.py

import mlx.core as mx
import mlx.nn as nn
import numpy as np
import librosa
from typing import Tuple, Optional, List
from .modules import RMSNorm, RotaryEmbedding, Transformer, MLP
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 MelBandRoformer(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, num_bands=60, dim_freqs_in=1025, sample_rate=44100, 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
        self.num_bands = num_bands
        
        # STFT params
        self.stft_n_fft = stft_n_fft
        self.stft_hop_length = stft_hop_length
        self.stft_win_length = stft_win_length
        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)

        # Create Mel Filter Bank
        try:
            mel_filters = librosa.filters.mel(sr=sample_rate, n_fft=stft_n_fft, n_mels=num_bands)
            mel_filters[0, 0] = 1.0
            mel_filters[-1, -1] = 1.0
            mel_filters = mx.array(mel_filters) # (n_mels, n_fft//2 + 1)
        except Exception as e:
            print(f"Error creating mel filters: {e}")
            raise e
            
        # Create boolean mask for bands
        freqs_per_band = (mel_filters > 0).astype(mx.float32) # (Bands, Freqs)
        self.freqs_per_band = freqs_per_band
        
        # Calculate number of frequencies per band for BandSplit
        num_freqs_per_band = mx.sum(freqs_per_band, axis=1).astype(mx.int32)
        
        # Prepare frequency indices for gathering
        freq_indices = []
        for b in range(num_bands):
            # Get indices where mask is > 0
            # helper to get indices from boolean mask in mlx/numpy
            indices_for_band = np.where(np.array(freqs_per_band[b]) > 0)[0]
            freq_indices.append(indices_for_band)
        
        self.freq_indices = mx.array(np.concatenate(freq_indices))
        
        # For reconstruction normalization (overlap count)
        self.num_bands_per_freq = mx.sum(freqs_per_band, axis=0) # (Freqs,)

        # Define dim_inputs for BandSplit
        dim_inputs = tuple(2 * f.item() * self.audio_channels for f in num_freqs_per_band)
        
        self.band_split, rope = BandSplit(dim, dim_inputs), 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, num_freqs_per_band.tolist(), depth=mask_estimator_depth)
            for _ in range(num_stems)
        ])
        
        self.fast_compute_masks = mx.compile(self._compute_masks)

    def _compute_masks(self, x):
        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
            
            # Time Transformer
            x = tt(x.transpose(0, 2, 1, 3).reshape(Bc * Fc, Tc, Dc)).reshape(Bc, Fc, Tc, Dc).transpose(0, 2, 1, 3)
            
            # Freq Transformer
            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)

    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) # (B*S, F_stft, Tf, 2)
        
        if self.stft_normalized:
            stft_repr = stft_repr * (self.stft_n_fft ** -0.5)
        
        if self.zero_dc:
            F_stft = stft_repr.shape[1]
            mask_list = [mx.zeros((1, 1, 1)), mx.ones((F_stft - 1, 1, 1))]
            mask = mx.concatenate(mask_list, axis=0) # (F, 1, 1)
            stft_repr = stft_repr * mask
            
        F_stft = stft_repr.shape[1]
        Tf = stft_repr.shape[2]
        
        # 1. Expand Stereo: (B, S, F, Tf, 2)
        stft_repr = stft_repr.reshape(B, S, F_stft, Tf, 2)
        
        # 2. Gather frequencies for all bands [Band1Freqs, Band2Freqs, ...]
        # Using take along axis 2 (Freq dim)
        stft_gathered = mx.take(stft_repr, self.freq_indices, axis=2)
        
        # Reshape for BandSplit: (B, Tf, TotalFreqs * S * 2)
        x_backbone = stft_gathered.transpose(0, 3, 2, 1, 4)
        x_backbone = x_backbone.reshape(B, Tf, -1)
        
        # Compute Masks: (B, NumStems, Tf, TotalFreqs * S * 2)
        masks = self.fast_compute_masks(x_backbone)
        
        # Reshape to (B, NumStems, Tf, TotalFreqs, S, 2)
        masks = masks.reshape(B, self.num_stems, Tf, -1, S, 2)
        
        # Scatter add to sum overlapping masks
        masks_summed = mx.zeros((B, self.num_stems, Tf, F_stft, S, 2))
        masks_summed = masks_summed.at[:, :, :, self.freq_indices, :, :].add(masks)
        
        # Normalize by overlap count
        denom = self.num_bands_per_freq[None, None, None, :, None, None]
        denom = mx.maximum(denom, 1e-8)
        
        masks_averaged = masks_summed / denom
        
        # Apply masks
        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_averaged[..., 0], masks_averaged[..., 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)
        
        # Prepare for iSTFT
        stft_masked = stft_masked.transpose(0, 1, 4, 3, 2, 5)
        stft_masked_flat = stft_masked.reshape(B * self.num_stems * S, F_stft, 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)