Move to the Gradio SDK, run inference in-process, and return every stem.

Dockerfile

@@ -1,24 +0,0 @@
-FROM python:3.10-slim
-
-WORKDIR /app
-
-RUN apt-get update && \
-    apt-get install -y --no-install-recommends git gcc libc6-dev ffmpeg && \
-    rm -rf /var/lib/apt/lists/*
-
-COPY requirements.txt /app/requirements.txt
-RUN pip install --no-cache-dir -r /app/requirements.txt
-
-# Pre-download model weights at build time so first inference is fast
-RUN bs-roformer-download \
-    --model roformer-model-bs-roformer-sw-by-jarredou \
-    --output-dir /app/models
-
-COPY . /app
-
-ENV PORT=7860
-EXPOSE 7860
-ENV PYTHONUNBUFFERED=1
-ENV PYTHONIOENCODING=UTF-8
-
-CMD ["python", "-u", "app.py"]

README.md

@@ -3,12 +3,13 @@ title: BS-RoFormer
 emoji: ✂️
 colorFrom: red
 colorTo: purple
-sdk: docker
+sdk: gradio
 sdk_version: 6.24.0
+python_version: "3.10"
 app_file: app.py
 pinned: false
 license: mit
-short_description: Band-split RoPE transformer stem separation
+short_description: Band-split RoPE transformer, six stems at once
 tags:
   - category:separation
   - subcategory:stem-separation
@@ -18,17 +19,13 @@ tags:
   - vocals
   - drums
   - bass
-app_port: 7860
 ---
 
-# BS-RoFormer — HARP Endpoint
+# BS-RoFormer
 
-Deploys [BS-RoFormer](https://github.com/lucidrains/BS-RoFormer) as a [PyHARP](https://github.com/TEAMuP-dev/pyharp) endpoint.
+Deploys [BS-RoFormer](https://github.com/lucidrains/BS-RoFormer) as a
+[PyHARP](https://github.com/TEAMuP-dev/pyharp) endpoint.
 
-## What it does
-
-Separates audio into stems (vocals, drums, bass, guitar, piano, other, instrumental) using the Band-Split RoPE Transformer — state-of-the-art music source separation from ByteDance AI Labs.
-
-## License
-
-MIT
+Separates a mixture into vocals, drums, bass, guitar, piano, and other with the SW checkpoint,
+which estimates all six in one pass, and returns each as its own track along with an instrumental
+(the mixture without its vocals). The output keeps the input's sample rate and channel count.

app.py

@@ -1,99 +1,133 @@
-import os
-import shutil
-import subprocess
+"""
+BS-RoFormer: separates a mixture into six stems with the Band-Split RoPE Transformer.
+
+The SW checkpoint estimates vocals, drums, bass, guitar, piano, and other in a single pass, so
+every stem is returned at once rather than one chosen ahead of time.
+"""
+
+import spaces  # imported before torch, as ZeroGPU requires
+
 import tempfile
 from pathlib import Path
 
 import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+import torchaudio
+import yaml
+from bs_roformer.download import default_models_dir, ensure_model_assets
+from bs_roformer.inference import SafeLoaderWithTuple
+from bs_roformer.utils import demix_track, get_model_from_config
+from filelock import FileLock
+from ml_collections import ConfigDict
 from pyharp import ModelCard, build_endpoint
 from pyharp.tags import Subcategory
 
-MODELS_DIR = Path("/app/models")
 MODEL_SLUG = "roformer-model-bs-roformer-sw-by-jarredou"
-MODEL_DIR = MODELS_DIR / MODEL_SLUG
+# The checkpoint was trained on 44.1 kHz stereo, whatever rate the input arrives at
+MODEL_SR = 44100
+# Presentation order. The checkpoint emits them as bass, drums, other, vocals, guitar, piano.
+STEMS = ["vocals", "drums", "bass", "guitar", "piano", "other"]
 
-STEMS = ["vocals", "drums", "bass", "guitar", "piano", "other", "instrumental"]
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-def get_model_paths():
-    configs = list(MODEL_DIR.glob("*.yaml"))
-    checkpoints = list(MODEL_DIR.glob("*.ckpt"))
-    if not configs or not checkpoints:
-        raise FileNotFoundError(f"Model files not found in {MODEL_DIR}")
-    return str(configs[0]), str(checkpoints[0])
+# pyharp's worker process imports this file as well, so both processes reach this point at
+# startup. The downloader writes straight to the final path, which would let the second one
+# find a partial checkpoint and load it, so only one downloads at a time.
+models_dir = default_models_dir()
+models_dir.mkdir(parents=True, exist_ok=True)
+with FileLock(str(models_dir / ".download.lock")):
+    checkpoint_path, config_path = ensure_model_assets(MODEL_SLUG, models_dir)
+
+with open(config_path) as f:
+    config = ConfigDict(yaml.load(f, Loader=SafeLoaderWithTuple))
+
+model = get_model_from_config("bs_roformer", config)
+model.load_state_dict(torch.load(checkpoint_path, map_location="cpu"))
+model.eval()
+model.to(device)
 
 model_card = ModelCard(
     name="BS-RoFormer",
     description=(
-        "Separate audio into stems (vocals, drums, bass, guitar, piano, other) "
-        "using the Band-Split RoPE Transformer. Repository: "
-        "https://github.com/lucidrains/BS-RoFormer Paper: "
-        "https://arxiv.org/abs/2309.02612"
+        "Separates a mixture into vocals, drums, bass, guitar, piano, and other with the "
+        "Band-Split RoPE Transformer, returning every stem along with an instrumental, which is "
+        "the mixture without its vocals. Repository: https://github.com/lucidrains/BS-RoFormer "
+        "Paper: https://arxiv.org/abs/2309.02612"
     ),
     author="Wei-Tsung Lu, Ju-Chiang Wang, Qiuqiang Kong, Yun-Ning Hung (ByteDance)",
-    tags=[Subcategory.STEM_SEPARATION, 'stems', 'vocals', 'drums', 'bass'],
+    tags=[Subcategory.STEM_SEPARATION, "stems", "vocals", "drums", "bass"],
 )
 
-def process_fn(input_audio_path: str, stem: str) -> str:
-    print(f"Separating stem: {stem}...", flush=True)
 
-    input_dir = Path(tempfile.mkdtemp())
-    output_dir = Path(tempfile.mkdtemp())
+@spaces.GPU(duration=120)
+@torch.inference_mode()
+def process_fn(input_audio_path: str):
+    mix, sr = sf.read(input_audio_path, always_2d=True)
+    num_samples, num_channels = mix.shape
+
+    if num_channels > 2:
+        raise gr.Error(f"Expected mono or stereo audio, but the input has {num_channels} channels.")
 
-    input_path = Path(input_audio_path)
-    shutil.copy(input_path, input_dir / input_path.name)
+    # The model only accepts stereo, so a mono input is duplicated and the stems folded back
+    stereo = np.repeat(mix, 2, axis=1) if num_channels == 1 else mix
+    mixture = torch.from_numpy(np.ascontiguousarray(stereo.T, dtype=np.float32))
 
-    config_path, model_path = get_model_paths()
-    print(f"Config: {config_path}", flush=True)
-    print(f"Model: {model_path}", flush=True)
-    print(f"Input dir contents: {list(input_dir.iterdir())}", flush=True)
+    if sr != MODEL_SR:
+        mixture = torchaudio.functional.resample(mixture, sr, MODEL_SR)
 
-    result = subprocess.run([
-        "bs-roformer-infer",
-        "--config_path", config_path,
-        "--model_path", model_path,
-        "--input_folder", str(input_dir),
-        "--store_dir", str(output_dir),
-    ], capture_output=True, text=True)
+    estimates, _ = demix_track(config, model, mixture, device)
 
-    print(f"STDOUT:\n{result.stdout}", flush=True)
-    print(f"STDERR:\n{result.stderr}", flush=True)
+    # The library derives its own instrumental from the first instrument, which for this
+    # checkpoint is the bass rather than the vocals, so it is computed here instead
+    estimates["instrumental"] = mixture.numpy() - estimates["vocals"]
 
-    if result.returncode != 0:
-        raise RuntimeError(f"bs-roformer-infer failed:\n{result.stderr}")
+    out_dir = Path(tempfile.mkdtemp())
+    stem = Path(input_audio_path).stem
+    output_paths = []
 
-    stem_files = list(output_dir.glob(f"*_{stem}.wav"))
-    if not stem_files:
-        available = [f.name for f in output_dir.glob("*.wav")]
-        raise ValueError(f"Stem '{stem}' not found. Available: {available}")
+    for name in STEMS + ["instrumental"]:
+        audio = torch.from_numpy(np.ascontiguousarray(estimates[name], dtype=np.float32))
 
-    output_path = tempfile.mktemp(suffix=".wav")
-    shutil.copy(stem_files[0], output_path)
-    print("Done.", flush=True)
-    return output_path
+        if sr != MODEL_SR:
+            audio = torchaudio.functional.resample(audio, MODEL_SR, sr)
 
+        # Resampling can shift the length by a sample or two
+        audio = audio[:, :num_samples]
+        if audio.shape[1] < num_samples:
+            audio = torch.nn.functional.pad(audio, (0, num_samples - audio.shape[1]))
+
+        audio = audio.mean(dim=0, keepdim=True) if num_channels == 1 else audio
+
+        path = out_dir / f"{stem}_{name}.wav"
+        sf.write(path, audio.numpy().T, sr, subtype="FLOAT")
+        output_paths.append(str(path))
+
+    return tuple(output_paths)
+
+
+# The processing worker imports this file, so the app must not be built there
 if __name__ == "__main__":
     with gr.Blocks() as demo:
         input_components = [
-            gr.Audio(
-                type="filepath",
-                label="Input Audio",
-            ).set_info("The mixture to separate."),
-            gr.Dropdown(
-                choices=STEMS,
-                value="vocals",
-                label="Stem to Extract",
-                info=(
-                    "Which part to pull out. Each stem uses its own trained model, so "
-                    "only one comes back per run."
-                ),
+            gr.Audio(type="filepath", label="Input Audio").set_info(
+                "The mixture to separate, in mono or stereo."
             ),
         ]
         output_components = [
-            gr.Audio(
-                type="filepath",
-                label="Separated Stem",
-            ).set_info("Extracted stem audio."),
+            gr.Audio(type="filepath", label=f"{name.capitalize()}").set_info(info)
+            for name, info in [
+                ("vocals", "The vocals."),
+                ("drums", "The drums and percussion."),
+                ("bass", "The bass."),
+                ("guitar", "The guitar."),
+                ("piano", "The piano."),
+                ("other", "Everything not assigned to another stem."),
+                ("instrumental", "The mixture without its vocals."),
+            ]
         ]
+
         app = build_endpoint(
             model_card=model_card,
             input_components=input_components,
@@ -101,5 +135,4 @@ if __name__ == "__main__":
             process_fn=process_fn,
         )
 
-    print("Launching Gradio...", flush=True)
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)
+    demo.queue().launch(show_error=True, pwa=True)

packages.txt

@@ -0,0 +1 @@
+ffmpeg

requirements.txt

@@ -1,3 +1,9 @@
-git+https://github.com/TEAMuP-dev/pyharp.git@v0.4.1
-bs-roformer-infer
+# pyharp v0.4.1
+git+https://github.com/TEAMuP-dev/pyharp.git@9a31acb35563c0bcf5a93f968d8285c7c3180aba
 spaces>=0.51.1
+# ZeroGPU requires torch 2.8 or newer
+torch==2.8.0
+torchaudio==2.8.0
+bs-roformer-infer==0.1.5
+soundfile
+filelock