Build all three models in the worker from a Hub mirror, declare the format, and pin requirements

Builds htdemucs, mdx_extra, and mdx_extra_q once, in the worker, from checkpoints in teamup-tech/Demucs-weights rather than inside every GPU call from Meta's servers, declares 44.1 kHz stereo, rewrites the description, removes a dead branch, writes to a temporary directory, drops the unused ffmpeg and scipy requirements, and pins the rest.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

README.md

@@ -15,6 +15,8 @@ tags:
   - subcategory:stem-separation
   - input:audio
   - output:audio
+  - sample-rate:44100
+  - channels:stereo
   - stems
 ---
 

app.py

@@ -1,57 +1,59 @@
 import spaces
 from pyharp.core import ModelCard, build_endpoint
 from pyharp.media.audio import save_audio
-from pyharp.labels import LabelList
-from pyharp.tags import Subcategory
+from pyharp.tags import Channels, SampleRate, Subcategory
 
 from audiotools import AudioSignal
 from demucs import pretrained
 from demucs.apply import apply_model
+from huggingface_hub import snapshot_download
 
+import contextlib
 import gradio as gr
 import soundfile as sf
+import tempfile
 import torch
 from pathlib import Path
-import numpy as np
 
 # ModelCard
 model_card = ModelCard(
     name="Demucs",
     description=(
-        "Separates a full music track into four individual audio stems: Drums, "
-        "Bass, Vocals, Other-instrumental using the Demucs model. Available "
-        "models are - htdemucs: Best overall fidelity and natural sound, "
-        "mdx_extra: Balanced trade-off, faster processing with a small drop in "
-        "quality, mdx_extra_q: Fastest & lightest, noticeable quality loss, "
-        "ideal for quick previews. Repository: "
+        "Separates a full music track into four stems: drums, bass, vocals, and "
+        "everything else. Three models are available. htdemucs, the Hybrid Transformer "
+        "Demucs, gives the best overall fidelity. mdx_extra is faster with a small drop in "
+        "quality, and mdx_extra_q is mdx_extra quantized, the lightest of the three, with a "
+        "noticeable loss of quality. Separation runs at 44.1 kHz in stereo, and the stems "
+        "come back at the rate and channel count of the input. Repository: "
         "https://github.com/facebookresearch/demucs Paper: "
         "https://arxiv.org/abs/2111.03600"
     ),
     author="Simon Rouard, Francisco Massa, Alexandre Défossez",
-    tags=[Subcategory.STEM_SEPARATION, 'stems']
+    tags=[Subcategory.STEM_SEPARATION, SampleRate(44100), Channels(2), 'stems']
 )
 
 DEMUX_MODELS = ["htdemucs", "mdx_extra", "mdx_extra_q"]
 STEM_NAMES = ["Drums", "Bass", "Vocals", "Other"]
 LOSSLESS = {".wav", ".flac", ".aiff", ".aif"}
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Global model cache
+# Every model is built and placed in the worker, which runs process_fn and starts as the server
+# comes up, so that no GPU call builds one. demucs would download the checkpoints from Meta's own
+# server, so they are read from a pinned mirror on the Hub, linked where torch.hub looks first.
 LOADED_MODELS = {}
 
-# Each model's weights are fetched in the worker, which runs process_fn and starts as the server
-# comes up, rather than inside the GPU call the first time a request picks the model
 if __name__ != "__main__":
-    for _model_name in DEMUX_MODELS:
-        pretrained.get_model(_model_name)
+    _weights = Path(snapshot_download(
+        "teamup-tech/Demucs-weights", revision="e12e76af128e0935a0ad1734cb97974e5d27663d"
+    ))
+    _hub = Path(torch.hub.get_dir()) / "checkpoints"
+    _hub.mkdir(parents=True, exist_ok=True)
+    for _checkpoint in _weights.glob("*/*.th"):
+        with contextlib.suppress(FileExistsError):
+            (_hub / _checkpoint.name).symlink_to(_checkpoint)
 
-def get_cached_model(model_name: str):
-    if model_name not in LOADED_MODELS:
-        model = pretrained.get_model(model_name)
-        model.eval()
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        model.to(device)
-        LOADED_MODELS[model_name] = model
-    return LOADED_MODELS[model_name]
+    for _model_name in DEMUX_MODELS:
+        LOADED_MODELS[_model_name] = pretrained.get_model(_model_name).eval().to(DEVICE)
 
 # The mixture is separated in overlapping segments, so the cost grows with its length
 # PyHARP stops a job that runs past its timeout_s, 15 minutes by default, so a request whose
@@ -69,33 +71,21 @@ def gpu_seconds(audio_file_path: str, model_name: str) -> float:
 # Separation Logic
 @spaces.GPU(duration=gpu_seconds)
 def separate_all_stems(audio_file_path: str, model_name: str):
-    model = get_cached_model(model_name)
+    model = LOADED_MODELS[model_name]
 
     signal = AudioSignal(audio_file_path)
     original_sr = signal.sample_rate
     signal = signal.resample(44100)  # expects 44.1kHz
 
+    # The model separates in stereo, so a mono input is duplicated to feed it
     is_mono = signal.num_channels == 1
     if is_mono:
         signal.audio_data = signal.audio_data.repeat(1, 2, 1)  # [batch, channels, samples]
-        signal._num_channels = 2 
-
-    sr = signal.sample_rate  
-
-    # Ensure audio_data is a torch.Tensor
-    audio = signal.audio_data
-    if isinstance(audio, np.ndarray):
-        audio = torch.from_numpy(audio)
-
-    audio = audio.float()  # [channels, samples] or [channels, samples, ?]
-
-    # Remove extra trailing dimensions
-    if audio.ndim > 2:
-        audio = audio.squeeze()
 
-    # Final shape: [1, channels, samples]
-    waveform = audio.unsqueeze(0)
+    sr = signal.sample_rate
 
+    # [batch, channels, samples], with the one item in the batch
+    waveform = signal.audio_data.float()
 
     with torch.no_grad():
         stems_batch = apply_model(
@@ -124,7 +114,7 @@ def separate_all_stems(audio_file_path: str, model_name: str):
 def process_fn(audio_file_path, model_name):
     try:
         if model_name not in DEMUX_MODELS:
-            raise ValueError(f"Unsupported model selected: {model_name}")
+            raise gr.Error(f"Unsupported model selected: {model_name}")
         # Checked before the GPU is reserved, so a request the model cannot take spends none of
         # the caller's ZeroGPU allowance
         if not audio_file_path:
@@ -140,13 +130,13 @@ def process_fn(audio_file_path, model_name):
             )
         output_signals = separate_all_stems(audio_file_path, model_name)
         # A lossless input keeps its container. A lossy one comes back as WAV, since encoding the
-        # stems lossily again would expose the first encode's artifacts and add a second set.
+        # stems lossily again would expose the artifacts of the first encode and add a second set.
         input_ext = Path(audio_file_path).suffix.lower()
         extension = input_ext.lstrip(".") if input_ext in LOSSLESS else "wav"
+        output_dir = Path(tempfile.mkdtemp())
         outputs = []
         for stem_name, signal in zip(STEM_NAMES, output_signals):
-            filename = f"demucs_{model_name}_{stem_name.lower().replace(' ', '_')}.{extension}"
-            output_path = Path(filename)
+            output_path = output_dir / f"demucs_{model_name}_{stem_name.lower()}.{extension}"
             save_audio(signal, output_path)
             outputs.append(str(output_path))
         return tuple(outputs)

requirements.txt

@@ -4,15 +4,13 @@ git+https://github.com/TEAMuP-dev/pyharp.git@v0.4.2
 demucs==4.0.0
 dora-search==0.1.11
 einops==0.6.1
-julius>=0.2.3
-lameenc>=1.2
+julius==0.2.8
+lameenc==1.9.0
 openunmix==1.2.1
 
 # Audio & ML
 torch==2.8.0
 torchaudio==2.8.0
-ffmpeg
-soundfile
-scipy
-numpy<2
+soundfile==0.14.0
+numpy==1.26.4
 spaces>=0.51.1