Upload inference.py

inference.py

@@ -0,0 +1,574 @@
+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()