| import argparse |
| import json |
| import shutil |
| import tempfile |
| from pathlib import Path |
| from urllib.error import HTTPError, URLError |
| from urllib.request import Request, urlopen |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| import torch.nn.functional as F |
| from einops import pack, rearrange, unpack |
| from rotary_embedding_torch import RotaryEmbedding |
| from safetensors.torch import load_file |
| from torch import einsum, nn |
|
|
|
|
| def pack_one(tensor, pattern): |
| return pack([tensor], pattern) |
|
|
|
|
| def unpack_one(tensor, packed_shape, pattern): |
| return unpack(tensor, packed_shape, pattern)[0] |
|
|
|
|
| class Attend(nn.Module): |
| def __init__(self): |
| super().__init__() |
|
|
| def forward(self, q, k, v): |
| scale = q.shape[-1] ** -0.5 |
| sim = einsum('b h i d, b h j d -> b h i j', q, k) * scale |
| attn = sim.softmax(dim=-1) |
| return einsum('b h i j, b h j d -> b h i d', attn, v) |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.scale = dim ** 0.5 |
| self.gamma = nn.Parameter(torch.ones(dim)) |
|
|
| def forward(self, x): |
| return F.normalize(x, dim=-1) * self.scale * self.gamma |
|
|
|
|
| class FeedForward(nn.Module): |
| def __init__(self, dim, ff_mult): |
| super().__init__() |
| dim_inner = int(dim * ff_mult) |
| self.net = nn.Sequential( |
| RMSNorm(dim), |
| nn.Linear(dim, dim_inner), |
| nn.GELU(), |
| nn.Identity(), |
| nn.Linear(dim_inner, dim), |
| nn.Identity(), |
| ) |
|
|
| def forward(self, x): |
| return self.net(x) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, dim, heads, dim_head, rotary_embed): |
| super().__init__() |
| self.heads = heads |
| dim_inner = heads * dim_head |
| self.rotary_embed = rotary_embed |
| self.attend = Attend() |
| self.norm = RMSNorm(dim) |
| self.to_qkv = nn.Linear(dim, dim_inner * 3, bias=False) |
| self.to_gates = nn.Linear(dim, heads) |
| self.to_out = nn.Sequential( |
| nn.Linear(dim_inner, dim, bias=False), |
| nn.Identity(), |
| ) |
|
|
| def forward(self, x): |
| x = self.norm(x) |
| q, k, v = rearrange( |
| self.to_qkv(x), |
| 'b n (qkv h d) -> qkv b h n d', |
| qkv=3, |
| h=self.heads, |
| ) |
|
|
| q = self.rotary_embed.rotate_queries_or_keys(q) |
| k = self.rotary_embed.rotate_queries_or_keys(k) |
|
|
| out = self.attend(q, k, v) |
| gates = self.to_gates(x) |
| out = out * rearrange(gates, 'b n h -> b h n 1').sigmoid() |
| out = rearrange(out, 'b h n d -> b n (h d)') |
| return self.to_out(out) |
|
|
|
|
| class Transformer(nn.Module): |
| def __init__(self, depth, dim, heads, dim_head, ff_mult, rotary_embed): |
| super().__init__() |
| self.layers = nn.ModuleList([]) |
|
|
| for _ in range(depth): |
| self.layers.append( |
| nn.ModuleList( |
| [ |
| Attention( |
| dim=dim, |
| heads=heads, |
| dim_head=dim_head, |
| rotary_embed=rotary_embed, |
| ), |
| FeedForward(dim=dim, ff_mult=ff_mult), |
| ] |
| ) |
| ) |
|
|
| def forward(self, x): |
| for attn, ff in self.layers: |
| x = attn(x) + x |
| x = ff(x) + x |
| return x |
|
|
|
|
| class BandSplit(nn.Module): |
| def __init__(self, dim_inputs, feature_dim): |
| super().__init__() |
| self.dim_inputs = dim_inputs |
| self.to_features = nn.ModuleList( |
| [nn.Sequential(nn.Linear(dim_in, feature_dim)) for dim_in in dim_inputs] |
| ) |
|
|
| def forward(self, x): |
| splits = x.split(self.dim_inputs, dim=-1) |
| features = [ |
| to_feature(split_input) |
| for split_input, to_feature in zip(splits, self.to_features) |
| ] |
| return torch.stack(features, dim=-2) |
|
|
|
|
| def MLP(dim_in, dim_out, dim_hidden, depth=1): |
| dims = (dim_in, *((dim_hidden,) * (depth - 1)), dim_out) |
|
|
| layers = [] |
| for index, (layer_dim_in, layer_dim_out) in enumerate(zip(dims[:-1], dims[1:])): |
| is_last = index == len(dims) - 2 |
| layers.append(nn.Linear(layer_dim_in, layer_dim_out)) |
| if not is_last: |
| layers.append(nn.Tanh()) |
|
|
| return nn.Sequential(*layers) |
|
|
|
|
| class MaskEstimator(nn.Module): |
| def __init__(self, dim_inputs, model_dim, depth, mlp_expansion_factor=4): |
| super().__init__() |
| dim_hidden = int(model_dim * mlp_expansion_factor) |
| self.to_freqs = nn.ModuleList( |
| [ |
| nn.Sequential( |
| MLP( |
| model_dim, |
| dim_in * 2, |
| dim_hidden=dim_hidden, |
| depth=depth, |
| ), |
| nn.GLU(dim=-1), |
| ) |
| for dim_in in dim_inputs |
| ] |
| ) |
|
|
| def forward(self, x): |
| outputs = [ |
| mlp(band_features) |
| for band_features, mlp in zip(x.unbind(dim=-2), self.to_freqs) |
| ] |
| return torch.cat(outputs, dim=-1) |
|
|
|
|
| class BSRoformer(nn.Module): |
| def __init__( |
| self, |
| *, |
| model_dim, |
| model_depth, |
| audio_channels, |
| num_stems, |
| time_transformer_depth, |
| freq_transformer_depth, |
| dim_head, |
| heads, |
| ff_mult, |
| stft_n_fft, |
| stft_hop_length, |
| stft_win_length, |
| stft_normalized, |
| mask_estimator_depth, |
| freq_range, |
| freqs_per_bands, |
| mask_mlp_expansion_factor=4, |
| ): |
| super().__init__() |
|
|
| self.audio_channels = audio_channels |
| self.num_stems = num_stems |
| self.layers = nn.ModuleList([]) |
|
|
| time_rotary_embed = RotaryEmbedding(dim=dim_head) |
| freq_rotary_embed = RotaryEmbedding(dim=dim_head) |
|
|
| for _ in range(model_depth): |
| self.layers.append( |
| nn.ModuleList( |
| [ |
| Transformer( |
| depth=time_transformer_depth, |
| dim=model_dim, |
| heads=heads, |
| dim_head=dim_head, |
| ff_mult=ff_mult, |
| rotary_embed=time_rotary_embed, |
| ), |
| Transformer( |
| depth=freq_transformer_depth, |
| dim=model_dim, |
| heads=heads, |
| dim_head=dim_head, |
| ff_mult=ff_mult, |
| rotary_embed=freq_rotary_embed, |
| ), |
| ] |
| ) |
| ) |
|
|
| self.final_norm = RMSNorm(model_dim) |
| self.stft_kwargs = dict( |
| n_fft=stft_n_fft, |
| hop_length=stft_hop_length, |
| win_length=stft_win_length, |
| normalized=stft_normalized, |
| ) |
| self.stft_window = torch.hann_window(stft_win_length) |
|
|
| freqs = stft_n_fft // 2 + 1 |
| min_freq, max_freq = (int(value) for value in freq_range) |
| if not 0 <= min_freq < max_freq <= freqs: |
| raise ValueError( |
| f'freq_range must satisfy 0 <= min < max <= {freqs}, got {(min_freq, max_freq)}' |
| ) |
| self.freq_slice = slice(min_freq, max_freq) |
| self.freq_pad = (min_freq, freqs - max_freq) |
|
|
| freqs_per_bands = tuple(int(band_size) for band_size in freqs_per_bands) |
| band_frequencies = max_freq - min_freq |
| if sum(freqs_per_bands) != band_frequencies: |
| raise ValueError( |
| f'freqs_per_bands must sum to {band_frequencies}, got {sum(freqs_per_bands)}' |
| ) |
|
|
| freqs_per_bands_with_complex = tuple( |
| 2 * band_size * self.audio_channels for band_size in freqs_per_bands |
| ) |
| self.band_split = BandSplit( |
| dim_inputs=freqs_per_bands_with_complex, |
| feature_dim=model_dim, |
| ) |
| self.mask_estimators = nn.ModuleList( |
| [ |
| MaskEstimator( |
| dim_inputs=freqs_per_bands_with_complex, |
| model_dim=model_dim, |
| depth=mask_estimator_depth, |
| mlp_expansion_factor=mask_mlp_expansion_factor, |
| ) |
| for _ in range(num_stems) |
| ] |
| ) |
|
|
| def forward(self, raw_audio): |
| if raw_audio.ndim == 2: |
| raw_audio = rearrange(raw_audio, 'b t -> b 1 t') |
|
|
| batch, channels, raw_audio_length = raw_audio.shape |
| if channels != self.audio_channels: |
| raise ValueError('audio channel count does not match the checkpoint architecture') |
|
|
| raw_audio, batch_audio_channel_packed_shape = pack_one(raw_audio, '* t') |
|
|
| stft_window = self.stft_window.to(device=raw_audio.device) |
|
|
| stft_repr = torch.stft( |
| raw_audio, |
| **self.stft_kwargs, |
| window=stft_window, |
| return_complex=True, |
| ) |
| stft_repr = torch.view_as_real(stft_repr) |
| stft_repr = unpack_one(stft_repr, batch_audio_channel_packed_shape, '* f t c') |
| stft_repr = stft_repr[:, :, self.freq_slice] |
| stft_repr = rearrange(stft_repr, 'b s f t c -> b (f s) t c') |
|
|
| x = rearrange(stft_repr, 'b f t c -> b t (f c)') |
| x = self.band_split(x) |
|
|
| for time_transformer, freq_transformer in self.layers: |
| x = rearrange(x, 'b t f d -> b f t d') |
| x, packed_shape = pack([x], '* t d') |
| x = time_transformer(x) |
| x, = unpack(x, packed_shape, '* t d') |
|
|
| x = rearrange(x, 'b f t d -> b t f d') |
| x, packed_shape = pack([x], '* f d') |
| x = freq_transformer(x) |
| x, = unpack(x, packed_shape, '* f d') |
|
|
| x = self.final_norm(x) |
|
|
| mask = torch.stack( |
| [mask_estimator(x) for mask_estimator in self.mask_estimators], |
| dim=1, |
| ) |
| mask = rearrange(mask, 'b n t (f c) -> b n f t c', c=2) |
|
|
| stft_repr = rearrange(stft_repr, 'b f t c -> b 1 f t c') |
| stft_repr = torch.view_as_complex(stft_repr) |
| mask = torch.view_as_complex(mask) |
| stft_repr = stft_repr * mask |
|
|
| stft_repr = rearrange( |
| stft_repr, |
| 'b n (f s) t -> (b n s) f t', |
| s=self.audio_channels, |
| ) |
| stft_repr = F.pad(stft_repr, (0, 0, *self.freq_pad)) |
|
|
| recon_audio = torch.istft( |
| stft_repr, |
| **self.stft_kwargs, |
| window=stft_window, |
| return_complex=False, |
| length=raw_audio_length, |
| ) |
|
|
| return rearrange( |
| recon_audio, |
| '(b n s) t -> b n s t', |
| b=batch, |
| s=self.audio_channels, |
| n=self.num_stems, |
| ) |
|
|
|
|
| INPUT_EXTENSIONS = {'.flac', '.wav', '.mp3'} |
| OUTPUT_FORMATS = {'wav', 'flac'} |
| DEFAULT_CONFIG_PATH = Path(__file__).with_name('config.json') |
| MODEL_CONFIG_URL = 'https://huggingface.co/tjpurdy/Piano-Separation-Model-small/resolve/main/config.json' |
| MODEL_CHECKPOINT_URL = 'https://huggingface.co/tjpurdy/Piano-Separation-Model-small/resolve/main/model.safetensors' |
| DOWNLOAD_TIMEOUT_SECONDS = 60 |
| MODEL_SAMPLE_RATE = 44100 |
| SEGMENT_SECONDS = 10 |
| DEFAULT_OVERLAP = 0.25 |
|
|
|
|
| def parse_output_format(value): |
| value = value.lower().lstrip('.') |
| if value not in OUTPUT_FORMATS: |
| raise argparse.ArgumentTypeError('output format must be wav or flac') |
| return value |
|
|
|
|
| def parse_overlap(value): |
| value = float(value) |
| if not 0 <= value < 1: |
| raise argparse.ArgumentTypeError('overlap must be in the range [0, 1)') |
| return value |
|
|
|
|
| def ensure_downloaded(file_path, url, description): |
| file_path = Path(file_path) |
| if file_path.exists(): |
| return file_path |
|
|
| file_path.parent.mkdir(parents=True, exist_ok=True) |
| temp_path = None |
| request = Request(url, headers={'User-Agent': 'inferencedownload/1.0'}) |
|
|
| try: |
| print(f'{description} not found at {file_path}, downloading from {url}') |
| with urlopen(request, timeout=DOWNLOAD_TIMEOUT_SECONDS) as response: |
| with tempfile.NamedTemporaryFile( |
| mode='wb', |
| delete=False, |
| dir=file_path.parent, |
| suffix='.download', |
| ) as temp_file: |
| temp_path = Path(temp_file.name) |
| shutil.copyfileobj(response, temp_file) |
|
|
| temp_path.replace(file_path) |
| print(f'Downloaded {description} to {file_path}') |
| return file_path |
| except (HTTPError, URLError, OSError) as exc: |
| if temp_path is not None and temp_path.exists(): |
| temp_path.unlink() |
| raise RuntimeError(f'Failed to download {description} from {url}: {exc}') from exc |
|
|
|
|
| def load_config(config_path): |
| config_path = ensure_downloaded(config_path, MODEL_CONFIG_URL, 'Model config') |
| with config_path.open('r', encoding='utf-8') as config_file: |
| return json.load(config_file) |
|
|
|
|
| def convert_audio(wav, from_sr, to_sr, channels): |
| if wav.ndim == 1: |
| wav = wav.unsqueeze(0) |
| if channels == 1: |
| wav = wav.mean(dim=0, keepdim=True) |
| elif wav.shape[0] == 1: |
| wav = wav.expand(channels, -1) |
| elif wav.shape[0] > channels: |
| wav = wav[:channels] |
| elif wav.shape[0] < channels: |
| raise ValueError('Audio has fewer channels than requested and is not mono.') |
| if from_sr == to_sr: |
| return wav |
|
|
| target_length = max(1, int(round(wav.shape[-1] * to_sr / from_sr))) |
| return F.interpolate( |
| wav.unsqueeze(0), |
| size=target_length, |
| mode='linear', |
| align_corners=False, |
| ).squeeze(0) |
|
|
|
|
| def load_separator(checkpoint_path, model_config, device): |
| model = BSRoformer(**model_config).eval().to(device) |
|
|
| checkpoint_path = Path(checkpoint_path) |
| checkpoint_was_missing = not checkpoint_path.exists() |
| checkpoint_path = ensure_downloaded( |
| checkpoint_path, |
| MODEL_CHECKPOINT_URL, |
| 'Model checkpoint', |
| ) |
| checkpoint_is_safetensors = checkpoint_was_missing or checkpoint_path.suffix == '.safetensors' |
| state = load_file(checkpoint_path) if checkpoint_is_safetensors else torch.load(checkpoint_path, map_location='cpu') |
| state = state.get('state', state) |
| model.load_state_dict({k[7:] if k.startswith('module.') else k: v for k, v in state.items()}) |
| return model |
|
|
|
|
| def list_audio_files(input_path): |
| input_path = Path(input_path) |
| if input_path.is_file(): |
| if input_path.suffix.lower() not in INPUT_EXTENSIONS: |
| raise ValueError(f'Input file is not a supported audio file: {input_path}') |
| return [input_path] |
|
|
| if not input_path.is_dir(): |
| raise FileNotFoundError( |
| f'Input path does not exist or is not a supported file/directory: {input_path}' |
| ) |
|
|
| files = sorted( |
| path |
| for path in input_path.rglob('*') |
| if path.is_file() and path.suffix.lower() in INPUT_EXTENSIONS |
| ) |
| duplicates = {} |
| for path in files: |
| duplicates.setdefault(path.stem, []).append(path) |
| duplicates = {stem: paths for stem, paths in duplicates.items() if len(paths) > 1} |
| if duplicates: |
| details = '\n'.join(f'{stem}: {", ".join(str(path) for path in paths)}' for stem, paths in sorted(duplicates.items())) |
| raise ValueError( |
| 'Multiple input files share the same name, so flat output filenames would collide:\n' + details |
| ) |
| return files |
|
|
|
|
| def run_model(model, mix, overlap): |
| length = mix.shape[-1] |
| segment = MODEL_SAMPLE_RATE * SEGMENT_SECONDS |
| stride = max(1, int(segment * (1 - overlap))) |
| weight = torch.cat(( |
| torch.arange(1, segment // 2 + 1, device=mix.device), |
| torch.arange(segment - segment // 2, 0, -1, device=mix.device), |
| )).float() |
| estimate = None |
| sum_weight = torch.zeros(length, device=mix.device) |
|
|
| with torch.inference_mode(): |
| for start in range(0, length, stride): |
| chunk = mix[:, start:start + segment] |
| chunk_est = model(chunk[None])[0] |
| if estimate is None: |
| estimate = torch.zeros(*chunk_est.shape[:-1], length, device=mix.device) |
| chunk_weight = weight[:chunk.shape[-1]] |
| estimate[..., start:start + chunk.shape[-1]] += chunk_est * chunk_weight |
| sum_weight[start:start + chunk.shape[-1]] += chunk_weight |
|
|
| return estimate / sum_weight |
|
|
|
|
| def separate_file(model, file_path, device, overlap): |
| audio, sample_rate = sf.read(file_path, dtype='float32') |
| mix = torch.from_numpy(np.asarray(audio, np.float32)) |
| mix = mix.unsqueeze(0) if mix.ndim == 1 else mix.T |
| source_channels = mix.shape[0] |
| mix = convert_audio(mix.to(device), sample_rate, MODEL_SAMPLE_RATE, model.audio_channels) |
|
|
| mono = mix.mean(0) |
| mean = mono.mean() |
| std = mono.std().clamp_min(1e-8) |
| mix = (mix - mean) / std |
|
|
| estimate = run_model(model, mix, overlap)[0] * std + mean |
| estimate = convert_audio(estimate, MODEL_SAMPLE_RATE, sample_rate, source_channels) |
| return estimate.T.cpu().numpy(), sample_rate |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser(description='Music source separation inference') |
| parser.add_argument('--input_dir', type=str, required=True, help='Input audio file or directory containing audio files') |
| parser.add_argument( |
| '--output_dir', |
| type=str, |
| default=None, |
| help='Output directory to save separated audio (default: same location as input)', |
| ) |
| parser.add_argument('--config_path', type=str, default=str(DEFAULT_CONFIG_PATH), help='Path to model config JSON') |
| parser.add_argument('--checkpoint_path', type=str, default='./model.safetensors', help='Path to model checkpoint file') |
| parser.add_argument('--output_format', type=parse_output_format, default='wav', help='Output file format: wav or flac (default: wav)') |
| parser.add_argument('--overlap', type=parse_overlap, default=DEFAULT_OVERLAP, help='Chunk overlap ratio in [0, 1) (default: 0.25)') |
| return parser.parse_args() |
|
|
|
|
| def main(): |
| args = parse_args() |
| input_path = Path(args.input_dir) |
| model_config = load_config(args.config_path) |
| audio_files = list_audio_files(args.input_dir) |
| if not audio_files: |
| print(f'No supported audio files found in {args.input_dir}') |
| return |
|
|
| if args.output_dir is not None: |
| output_dir = Path(args.output_dir) |
| else: |
| output_dir = input_path.parent if input_path.is_file() else input_path |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| if device.type == 'cpu': |
| print('WARNING, using CPU') |
|
|
| model = load_separator(args.checkpoint_path, model_config, device) |
| print(f'Found {len(audio_files)} audio file(s) from {args.input_dir}') |
|
|
| for file_path in audio_files: |
| print(f'Processing {file_path}') |
| estimate, sample_rate = separate_file(model, file_path, device, args.overlap) |
| save_path = output_dir / f'{file_path.stem}_Piano.{args.output_format}' |
| sf.write(save_path, estimate, sample_rate) |
| print(f'Saved {save_path}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
