infer.py

"""
Pure numpy + onnxruntime reference implementation for the HT-Demucs FT
drums specialist. NO TORCH at inference — works in any environment with
numpy, soundfile, and onnxruntime.

Usage:
    python infer.py input.mp3 out_dir/
    # writes out_dir/drums.wav (the meaningful stem)

Or as a library:
    import infer
    drums = infer.separate_drums("song.mp3", "htdemucs_ft_drums.onnx")
    # drums: numpy array (channels, samples) at 44.1 kHz

This script was hand-rolled to be readable and portable rather than maximally
fast. For production deployment you'll want to:
  - Use IOBinding to skip the numpy<->ORT copies
  - Use GraphOptimizationLevel.ORT_ENABLE_ALL
  - Add CoreMLExecutionProvider / DmlExecutionProvider / CUDAExecutionProvider
    depending on platform (see https://onnxruntime.ai/docs/execution-providers/)
"""
from __future__ import annotations

import argparse
import sys
import time
from pathlib import Path

import numpy as np
import onnxruntime as ort
import soundfile as sf

# These constants are baked into the exported graph. If you re-export with
# a different segment length, update them here.
SAMPLE_RATE = 44100
SEGMENT_S = 7.8
N_SAMPLES = int(SEGMENT_S * SAMPLE_RATE)  # 343,980
N_CHANNELS = 2
SOURCES = ["drums", "bass", "other", "vocals"]  # the bag's output order
SPECIALIST_STEM = "drums"  # which stem is meaningfully predicted by this model
DEFAULT_ONNX = Path(__file__).resolve().parent / "htdemucs_ft_drums.onnx"


def _make_transition_window(segment: int, overlap_frac: float = 0.25) -> np.ndarray:
    transition = int(segment * overlap_frac)
    window = np.ones(segment, dtype=np.float32)
    fade = np.linspace(0, 1, transition, dtype=np.float32)
    window[:transition] = fade
    window[-transition:] = fade[::-1]
    return window


def _load_session(onnx_path: Path,
                  providers: list[str] | None = None) -> ort.InferenceSession:
    if providers is None:
        # CPU-only default. To enable CoreML on macOS, pass:
        #   providers=["CoreMLExecutionProvider", "CPUExecutionProvider"]
        # (first-time CoreML compile of this 24k-node graph is SLOW; expect
        # 5+ minutes the first run. Subsequent loads are fast.)
        providers = ["CPUExecutionProvider"]
    return ort.InferenceSession(str(onnx_path), providers=providers)


def separate(mix: np.ndarray, sample_rate: int,
             onnx_path: Path = DEFAULT_ONNX,
             providers: list[str] | None = None,
             verbose: bool = True) -> np.ndarray:
    """Run chunked overlap-add separation on a full-length mix.

    Args:
        mix: (channels, samples) float32 in [-1, 1]. Must be 2-channel stereo
             at SAMPLE_RATE. Resample first with `soxr`, `librosa`, or `ffmpeg`
             if needed.
        sample_rate: must equal SAMPLE_RATE.
        onnx_path: path to htdemucs_ft_drums.onnx (defaults to alongside this
                   script).
        providers: list of onnxruntime execution providers; defaults to CPU.
        verbose: print per-chunk progress.

    Returns:
        (n_sources, channels, samples) float32 array. Only the row indexed by
        SOURCES.index(SPECIALIST_STEM) is meaningfully predicted — the other
        three rows are weakly-predicted by-products of the drum specialist
        and should not be used in production. For high-quality bass / vocals
        / other, ship the respective specialist ONNX file instead.
    """
    if sample_rate != SAMPLE_RATE:
        raise ValueError(
            f"This model is bound to {SAMPLE_RATE} Hz; got {sample_rate}. "
            "Resample your input first.")
    if mix.ndim != 2 or mix.shape[0] != N_CHANNELS:
        raise ValueError(f"Expected (2, samples) input, got {mix.shape}")

    sess = _load_session(onnx_path, providers)

    total_len = mix.shape[1]
    overlap = N_SAMPLES // 4
    stride = N_SAMPLES - overlap
    n_chunks = max(1, (total_len + stride - 1) // stride)

    if verbose:
        print(f"  input:    {total_len:,} samples ({total_len / sample_rate:.1f}s)")
        print(f"  segment:  {N_SAMPLES:,} samples ({SEGMENT_S}s)")
        print(f"  chunks:   {n_chunks}, stride {stride / sample_rate:.2f}s, "
              f"overlap {overlap / sample_rate:.2f}s")
        print(f"  provider: {sess.get_providers()[0]}")

    window = _make_transition_window(N_SAMPLES)
    out = np.zeros((len(SOURCES), N_CHANNELS, total_len), dtype=np.float32)
    weight = np.zeros(total_len, dtype=np.float32)

    t0 = time.perf_counter()
    for i in range(n_chunks):
        start = i * stride
        end = min(start + N_SAMPLES, total_len)
        chunk = mix[:, start:end]
        if chunk.shape[1] < N_SAMPLES:
            chunk = np.pad(chunk, ((0, 0), (0, N_SAMPLES - chunk.shape[1])),
                           mode="constant")
        x = chunk[np.newaxis, ...].astype(np.float32)  # (1, 2, N)
        stems = sess.run(["stems"], {"mix": x})[0][0]  # (4, 2, N)
        chunk_len = end - start
        w = window[:chunk_len]
        out[:, :, start:end] += stems[:, :, :chunk_len] * w
        weight[start:end] += w
        if verbose:
            print(f"    chunk {i+1}/{n_chunks}: "
                  f"{time.perf_counter() - t0:.1f}s elapsed")

    weight = np.maximum(weight, 1e-8)
    out /= weight
    if verbose:
        rtf = (time.perf_counter() - t0) / (total_len / sample_rate)
        print(f"  total:    {time.perf_counter() - t0:.2f}s (RTF {rtf:.2f})")
    return out


def separate_drums(input_path: str, onnx_path: Path = DEFAULT_ONNX,
                   providers: list[str] | None = None) -> np.ndarray:
    """Convenience: load audio, run separation, return only the drums stem.

    Returns: (channels, samples) float32 at 44.1 kHz.
    """
    audio, sr = sf.read(input_path, dtype="float32", always_2d=True)
    audio = audio.T
    if audio.shape[0] == 1:
        audio = np.tile(audio, (2, 1))
    elif audio.shape[0] > 2:
        audio = audio[:2]
    stems = separate(audio, sr, onnx_path=onnx_path, providers=providers)
    return stems[SOURCES.index(SPECIALIST_STEM)]


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("input", type=Path, help="Input audio file (wav/mp3/flac).")
    ap.add_argument("out_dir", type=Path, help="Output directory for .wav stems.")
    ap.add_argument("--onnx", type=Path, default=DEFAULT_ONNX,
                    help=f"Path to the .onnx file (default: {DEFAULT_ONNX})")
    ap.add_argument("--providers", type=str, default="cpu",
                    choices=["cpu", "coreml", "cuda", "dml"],
                    help="Execution provider (cpu is default and most portable).")
    ap.add_argument("--write-all-stems", action="store_true",
                    help="Also write bass/other/vocals (by-products, low quality).")
    args = ap.parse_args()

    providers_map = {
        "cpu": ["CPUExecutionProvider"],
        "coreml": ["CoreMLExecutionProvider", "CPUExecutionProvider"],
        "cuda": ["CUDAExecutionProvider", "CPUExecutionProvider"],
        "dml": ["DmlExecutionProvider", "CPUExecutionProvider"],
    }
    providers = providers_map[args.providers]
    args.out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Loading {args.input} ...")
    audio, sr = sf.read(str(args.input), dtype="float32", always_2d=True)
    audio = audio.T
    if audio.shape[0] == 1:
        audio = np.tile(audio, (2, 1))
    elif audio.shape[0] > 2:
        audio = audio[:2]
    print(f"  shape {audio.shape}, sr {sr}")

    print(f"\nLoading {args.onnx} ({args.onnx.stat().st_size / 1e6:.1f} MB) ...")
    stems = separate(audio, sr, onnx_path=args.onnx, providers=providers)

    if args.write_all_stems:
        for i, src in enumerate(SOURCES):
            sf.write(str(args.out_dir / f"{src}.wav"), stems[i].T, sr)
            print(f"  wrote {args.out_dir / f'{src}.wav'}  "
                  f"{'*' if src == SPECIALIST_STEM else ''}")
        print("\n  * = meaningful prediction (specialist target)")
    else:
        drums = stems[SOURCES.index(SPECIALIST_STEM)]
        out_path = args.out_dir / "drums.wav"
        sf.write(str(out_path), drums.T, sr)
        print(f"  wrote {out_path}")


if __name__ == "__main__":
    main()