import subprocess
import sys
import tempfile
from pathlib import Path

# Install demucs without deps so its torchaudio<2.1 pin doesn't drag torch backwards
# and break the `spaces` import (which needs torch >= 2.1).
subprocess.run(
    [sys.executable, "-m", "pip", "install", "--no-deps", "git+https://github.com/facebookresearch/demucs"],
    check=True,
)

import spaces
import gradio as gr
from scipy.io.wavfile import write
from demucs.api import Separator, save_audio

# Preload weights to disk; CPU instantiation avoids hijacking CUDA at import time.
Separator(model="htdemucs", device="cpu")


@spaces.GPU
def inference(audio):
    sr, arr = audio
    in_path = Path(tempfile.mktemp(suffix=".wav"))
    write(in_path, sr, arr)

    separator = Separator(model="htdemucs")  # picks cuda inside the ZeroGPU fork
    _, stems = separator.separate_audio_file(in_path)
    vocals = stems["vocals"]
    no_vocals = sum(v for k, v in stems.items() if k != "vocals")

    voc_path = tempfile.mktemp(suffix=".wav")
    nv_path = tempfile.mktemp(suffix=".wav")
    save_audio(vocals, voc_path, samplerate=separator.samplerate)
    save_audio(no_vocals, nv_path, samplerate=separator.samplerate)
    return voc_path, nv_path


title = "Demucs Music Source Separation (v4)"
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1911.13254' target='_blank'>Music Source Separation in the Waveform Domain</a> | <a href='https://github.com/facebookresearch/demucs' target='_blank'>Github Repo</a> | <a href='https://www.thafx.com' target='_blank'>//THAFX</a></p>"

gr.Interface(
    inference,
    gr.Audio(type="numpy", label="Input"),
    [gr.Audio(type="filepath", label="Vocals"), gr.Audio(type="filepath", label="No Vocals / Instrumental")],
    title=title,
    article=article,
    examples=[["test.mp3"]],
    cache_examples=False,
).launch()