Spaces:
Running
Running
File size: 3,072 Bytes
f0bb30e bcb4375 f0bb30e bcb4375 f0bb30e bcb4375 f0bb30e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import os
import shutil
import subprocess
import tempfile
from pathlib import Path
import gradio as gr
from pyharp import ModelCard, build_endpoint
MODELS_DIR = Path("/app/models")
MODEL_SLUG = "roformer-model-bs-roformer-sw-by-jarredou"
MODEL_DIR = MODELS_DIR / MODEL_SLUG
STEMS = ["vocals", "drums", "bass", "guitar", "piano", "other", "instrumental"]
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])
model_card = ModelCard(
name="BS-RoFormer Source Separation",
description="Separate audio into stems (vocals, drums, bass, guitar, piano, other) using the Band-Split RoPE Transformer.",
author="Wei-Tsung Lu, Ju-Chiang Wang, Qiuqiang Kong, Yun-Ning Hung (ByteDance)",
tags=["source-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())
input_path = Path(input_audio_path)
shutil.copy(input_path, input_dir / input_path.name)
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)
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)
print(f"STDOUT:\n{result.stdout}", flush=True)
print(f"STDERR:\n{result.stderr}", flush=True)
if result.returncode != 0:
raise RuntimeError(f"bs-roformer-infer failed:\n{result.stderr}")
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}")
output_path = tempfile.mktemp(suffix=".wav")
shutil.copy(stem_files[0], output_path)
print("Done.", flush=True)
return output_path
with gr.Blocks() as demo:
input_components = [
gr.Audio(
type="filepath",
label="Input Audio",
),
gr.Dropdown(
choices=STEMS,
value="vocals",
label="Stem to Extract",
),
]
output_components = [
gr.Audio(
type="filepath",
label="Separated Stem",
).set_info("Extracted stem audio."),
]
app = build_endpoint(
model_card=model_card,
input_components=input_components,
output_components=output_components,
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)
|