updated

app.py

@@ -1,148 +1,105 @@
-# app.py
-
-import os
-import uuid
 import torch
 import torchaudio
-import soundfile as sf
 import gradio as gr
-
 from demucs import pretrained
 from demucs.apply import apply_model
+from audiotools import AudioSignal
+from pyharp.core import ModelCard, build_endpoint
+from pyharp.labels import LabelList
+from pyharp.media.audio import save_audio
 
-from pyharp.core import ModelCard, build_endpoint, extend_gradio
-from pyharp.labels import LabelList, AudioLabel
-
-# ───────────────────────────────────────────────────────────────
-# Spaces-compatible cache dirs (also works locally)
-# ───────────────────────────────────────────────────────────────
-os.environ["XDG_CACHE_HOME"] = "/tmp/.cache"
-os.environ["TORCH_HOME"] = "/tmp/torch"
-os.environ["HF_HOME"] = "/tmp/hf"
-os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf"
-
-for d in (os.environ["XDG_CACHE_HOME"], os.environ["TORCH_HOME"], os.environ["HF_HOME"], os.environ["MPLCONFIGDIR"]):
-    os.makedirs(d, exist_ok=True)
-
-# Ensure Gradio components support .harp_required()
-extend_gradio()
+# Supported models
+DEMUX_MODELS = ["mdx_extra_q", "mdx_extra", "htdemucs", "mdx_q"]
 
-# ───────────────────────────────────────────────────────────────
-# Constants
-# ───────────────────────────────────────────────────────────────
-DEMUCS_MODELS = ["mdx_extra_q", "mdx_extra", "htdemucs", "mdx_q"]
+# Mapping stem names to indexes
 STEM_CHOICES = {
     "Vocals": 3,
     "Drums": 0,
     "Bass": 1,
     "Other": 2,
-    "Instrumental (No Vocals)": "instrumental",
+    "Instrumental (No Vocals)": "instrumental"
 }
-DEMUCS_SR = 44100
-
-# ───────────────────────────────────────────────────────────────
-# Utilities
-# ───────────────────────────────────────────────────────────────
-def ensure_stereo(wave: torch.Tensor) -> torch.Tensor:
-    if wave.shape[0] == 1:
-        return wave.repeat(2, 1)
-    return wave[:2]
-
-def resample_if_needed(wave: torch.Tensor, sr: int, target_sr: int):
-    if sr == target_sr:
-        return wave
-    return torchaudio.functional.resample(wave, sr, target_sr)
-
-def save_wav(audio: torch.Tensor, sr: int, stem_name: str) -> str:
-    out_path = f"/tmp/{stem_name}_{uuid.uuid4().hex}.wav"
-    sf.write(out_path, audio.cpu().numpy().T, sr)
-    return out_path
-
-# ───────────────────────────────────────────────────────────────
-# Main processing function
-# ───────────────────────────────────────────────────────────────
-def process_fn(audio_path: str, model_name: str, stem_choice: str):
-    # Load and prepare audio
-    wave, sr = torchaudio.load(audio_path)
-    wave = ensure_stereo(wave.float())
-    wave = resample_if_needed(wave, sr, DEMUCS_SR)
-
-    # Load model
+
+
+def separate_stem(audio_file_path: str, model_name: str, stem_choice: str) -> AudioSignal:
     model = pretrained.get_model(model_name)
-    model.to("cpu").eval()
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model.to(device)
+    model.eval()
+
+    waveform, sr = torchaudio.load(audio_file_path)
+    is_mono = waveform.shape[0] == 1
+    if is_mono:
+        waveform = waveform.repeat(2, 1)
 
-    # Apply separation
     with torch.no_grad():
-        batch = wave.unsqueeze(0)
-        stems = apply_model(model, batch, overlap=0.2, shifts=1, split=True)[0]
+        stems_batch = apply_model(
+            model,
+            waveform.unsqueeze(0).to(device),
+            overlap=0.2,
+            shifts=1,
+            split=True
+        )
+
+    stems = stems_batch[0]
 
-    # Extract desired stem
     if stem_choice == "Instrumental (No Vocals)":
-        stem_audio = stems[0] + stems[1] + stems[2]
+        stem = stems[0] + stems[1] + stems[2]
     else:
-        stem_audio = stems[STEM_CHOICES[stem_choice]]
-
-    # Resample back if needed
-    if DEMUCS_SR != sr:
-        stem_audio = torchaudio.functional.resample(stem_audio, DEMUCS_SR, sr)
-
-    # Save to temp file
-    base_name = stem_choice.lower().replace(" ", "_").replace("(", "").replace(")", "")
-    output_path = save_wav(stem_audio, sr, base_name)
-
-    # Create simple label for full duration
-    label = AudioLabel(
-        t=0.0,
-        duration=stem_audio.shape[-1] / sr,
-        label=stem_choice,
-        amplitude=0.0,
-        color=AudioLabel.hex_color_to_int("#4CAF50"),
-    )
-    label_list = LabelList()
-    label_list.append(label)
+        stem_index = STEM_CHOICES[stem_choice]
+        stem = stems[stem_index]
+
+    if is_mono:
+        stem = stem.mean(dim=0, keepdim=True)
+
+    return AudioSignal(stem.cpu().numpy().astype('float32'), sample_rate=sr)
+
 
-    return output_path, label_list
+def process_fn_stem(audio_file_path: str, demucs_model: str, stem_choice: str):
+    """
+    PyHARP v3 process function:
+      - Separates the chosen stem using Demucs.
+      - Saves the stem as a .wav file.
+    """
+    stem_signal = separate_stem(audio_file_path, model_name=demucs_model, stem_choice=stem_choice)
+    stem_path = save_audio(stem_signal, f"{stem_choice.lower().replace(' ', '_')}.wav")
+    return stem_path
 
-# ───────────────────────────────────────────────────────────────
-# ModelCard for pyharp
-# ───────────────────────────────────────────────────────────────
+
+# Model Card
 model_card = ModelCard(
-    name="Demucs Stem Separator (CPU)",
-    description="Separate a mix into a chosen stem using Demucs (CPU-only).",
-    author="Your Name or Team",
-    tags=["demucs", "source-separation", "audio", "stem", "harp"],
+    name="Demucs Stem Separator",
+    description="Uses Demucs to separate a music track into a selected stem.",
+    author="Alexandre Défossez, Nicolas Usunier, Léon Bottou, Francis Bach",
+    tags=["demucs", "source-separation", "pyharp", "stems"]
 )
 
-# ───────────────────────────────────────────────────────────────
-# Gradio UI and pyharp integration
-# ───────────────────────────────────────────────────────────────
-with gr.Blocks(title=model_card.name) as demo:
-    # Define UI inputs
-    input_audio = gr.Audio(type="filepath", label="Input Audio").harp_required(True)
-    dropdown_model = gr.Dropdown(label="Demucs Model", choices=DEMUCS_MODELS, value="mdx_extra_q")
-    dropdown_stem = gr.Dropdown(label="Stem", choices=list(STEM_CHOICES.keys()), value="Vocals")
-
-    # Define UI outputs
-    output_audio = gr.Audio(type="filepath", label="Output Stem")
-    output_json = gr.JSON(label="Labels")
+# Gradio UI
+with gr.Blocks() as demo:
+
+    input_components = [
+        gr.Audio(type="filepath", label="Input Audio").harp_required(True),
+        gr.Dropdown(
+            label="Select Demucs Model",
+            choices=DEMUX_MODELS,
+            value="mdx_extra_q"
+        ),
+        gr.Dropdown(
+            label="Select Stem to Separate",
+            choices=list(STEM_CHOICES.keys()),
+            value="Vocals"
+        )
+    ]
+
+    output_components = [
+        gr.Audio(type="filepath", label="Separated Output"),
+    ]
 
-    # Build HARP-compatible endpoint inside the Blocks context
     app = build_endpoint(
         model_card=model_card,
-        input_components=[input_audio, dropdown_model, dropdown_stem],
-        output_components=[output_audio, output_json],
-        process_fn=process_fn,
+        input_components=input_components,
+        output_components=output_components,
+        process_fn=process_fn_stem
     )
 
-    # Add the control buttons for HARP
-    app["controls_button"]
-    app["controls_data"]
-    app["process_button"]
-    app["cancel_button"]
-
-# Queue and launch
-demo.queue()
-if __name__ == "__main__":
-   demo.launch(show_error=True, share=True)
-
+demo.queue().launch(share=True,show_error=True)