Spaces:

teamup-tech
/

Open-Unmix

Running on Zero

Add domain tags for the Home tab filters of HARP

8f89410 verified about 4 hours ago

8.46 kB

from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import contextlib
import tempfile
from pathlib import Path

import gradio as gr
import soundfile as sf
import torch
import torchaudio
from huggingface_hub import snapshot_download
from openunmix import predict, utils as umx_utils

from pyharp import *


model_card = ModelCard(
    name="Open-Unmix",
    description=(
        "Music source separation into vocals, drums, bass, and other "
        "instruments using deep neural networks. Repository: "
        "https://github.com/sigsep/open-unmix-pytorch Paper: "
        "https://doi.org/10.5281/zenodo.5069601 The code and the umxhq model "
        "are MIT, but the default umxl model carries CC-BY-NC-SA-4.0, so its "
        "results cannot be used commercially and must be shared alike."
    ),
    author="F.-R. Stöter, S. Uhlich, A. Liutkus, Y. Mitsufuji",
    tags=[Subcategory.STEM_SEPARATION, SampleRate(44100), Channels(2), 'stems', 'domain:music'],
)

MODELS = ["umxl", "umxhq"]
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the stems
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}

# Both separators are built and placed in the worker, which runs process_fn and starts with the
# server, so that no GPU call builds one. The Wiener filter settings are plain attributes of a
# separator, set for each request. openunmix downloads each checkpoint from Zenodo on first use,
# which is too slow for a Space to start, so the same files are read from a mirror on the Hub,
# linked where torch.hub looks before downloading.
SEPARATORS = {}

if __name__ != "__main__":
    _weights = Path(snapshot_download(
        "teamup-tech/Open-Unmix-weights", revision="2dca5d0f8b5692be0fa99c0d45c73fa7bff6d665"
    ))
    _hub = Path(torch.hub.get_dir()) / "checkpoints"
    _hub.mkdir(parents=True, exist_ok=True)
    for _checkpoint in _weights.glob("*/*.pth"):
        with contextlib.suppress(FileExistsError):
            (_hub / _checkpoint.name).symlink_to(_checkpoint)

    for _model_name in MODELS:
        _separator = umx_utils.load_separator(model_str_or_path=_model_name, device=DEVICE,
                                          pretrained=True)
        _separator.freeze()
        SEPARATORS[_model_name] = _separator.to(DEVICE)


# The whole mixture is separated, 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
# reservation would exceed that is refused rather than started. A shorter one is left to
# ZeroGPU, which admits it when the caller's own allowance covers it.
MAX_GPU_SECONDS = 900


def gpu_seconds(input_audio, model_name, niter, wiener_win_len):
    # Measured at about 2 s for 20 s of audio, 4 s for 60 s and 9 s for 300 s, upload included,
    # at any Wiener setting, padded. The separators are built when the worker starts, not inside
    # the call.
    return 5 + 0.04 * sf.info(input_audio).duration


def process_fn(input_audio, model_name, niter, wiener_win_len):
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if model_name not in MODELS:
        raise gr.Error(f"Unknown model: {model_name}")
    if not input_audio:
        raise gr.Error("Provide the mixture to separate.")
    try:
        duration = sf.info(input_audio).duration
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc
    if duration == 0:
        raise gr.Error("The audio file is empty.")
    if gpu_seconds(input_audio, model_name, niter, wiener_win_len) > MAX_GPU_SECONDS:
        raise gr.Error(
            f"The mixture is {duration:.0f} s long, more than one request can separate in the "
            f"{MAX_GPU_SECONDS // 60} minutes it may run. Trim it."
        )
    return separate(input_audio, model_name, niter, wiener_win_len)


# Inference mode as well as freeze(), since ZeroGPU recreates the parameters on the GPU for each
# call, and the copies would require gradients again
@spaces.GPU(duration=gpu_seconds)
@torch.inference_mode()
def separate(input_audio, model_name, niter, wiener_win_len):
    data, rate = sf.read(input_audio, dtype="float32", always_2d=True)
    audio = torch.from_numpy(data.T)

    # Open-Unmix separates at 44.1 kHz, so the session's rate is kept here and
    # the stems are resampled back to it before saving
    original_sr = rate
    original_channels = audio.shape[0]

    # A mono input is duplicated to feed the model, which separates in stereo. Two identical
    # channels leave the Wiener filter's spatial covariance singular, which past a couple of
    # iterations turns every stem to NaN, so the copy carries noise at -100 dBFS, below the
    # quietest step of a 16-bit file. The stems are averaged back to mono below.
    if audio.shape[0] == 1:
        noise = torch.randn(audio.shape, generator=torch.Generator().manual_seed(0))
        audio = torch.cat([audio, audio + 1e-5 * noise])

    separator = SEPARATORS[model_name]
    separator.niter = int(niter)
    separator.wiener_win_len = int(wiener_win_len)

    estimates = predict.separate(
        audio=audio,
        rate=rate,
        model_str_or_path=model_name,
        separator=separator,
        device=DEVICE
    )

    input_ext = Path(input_audio).suffix.lower()
    output_dir = Path(tempfile.mkdtemp())
    stems = {}
    for target in ["vocals", "drums", "bass", "other"]:
        target_tensor = estimates[target][0].cpu()

        if int(separator.sample_rate) != int(original_sr):
            target_tensor = torchaudio.functional.resample(
                target_tensor, int(separator.sample_rate), int(original_sr)
            )

        # A mono input was duplicated to feed the model, which separates in
        # stereo; giving back two identical-sounding channels would claim an
        # image the input never had
        if original_channels == 1:
            target_tensor = target_tensor.mean(dim=0, keepdim=True)

        stems[target] = target_tensor

    # Wiener filtering can push a stem past full scale, which a 16-bit file would flatten and a
    # FLAC file cannot hold at all, so when any stem exceeds it, all are scaled down together and
    # still sum to the mixture
    peak = max(float(stem.abs().max()) for stem in stems.values())
    output_files = {}
    for target, stem in stems.items():
        if peak > 1.0:
            stem = stem / peak
        out_file = output_dir / f"{target}{input_ext if input_ext in LOSSLESS else '.wav'}"
        sf.write(str(out_file), stem.numpy().T, int(original_sr))
        output_files[target] = str(out_file)

    return (
        output_files["vocals"],
        output_files["drums"],
        output_files["bass"],
        output_files["other"]
    )


if __name__ == "__main__":
    with gr.Blocks() as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                "The mixture to separate. Mono input is duplicated to stereo first."
            )),
            gr.Dropdown(choices=MODELS, value="umxl", label="Model", info="umxl separates best, but its results cannot be used commercially. umxhq is MIT licensed."),
            gr.Slider(minimum=0, maximum=5, step=1, value=1, label="Wiener Filter Iterations", info="Number of iterations for Wiener filtering. 0 means softmask."),
            gr.Slider(minimum=100, maximum=500, step=50, value=300, label="Wiener Window Length", info="Number of frames on which to apply filtering independently."),
        ]
        output_components = [
            gr.Audio(type="filepath", label="Vocals").set_info("The isolated vocal."),
            gr.Audio(type="filepath", label="Drums").set_info("The isolated drums."),
            gr.Audio(type="filepath", label="Bass").set_info("The isolated bass."),
            gr.Audio(type="filepath", label="Other").set_info((
                "Everything the other three stems did not take."
            )),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )

    demo.queue().launch(show_error=True, pwa=True)

	from __future__ import annotations

	import spaces # imported first, as ZeroGPU requires

	import contextlib
	import tempfile
	from pathlib import Path

	import gradio as gr
	import soundfile as sf
	import torch
	import torchaudio
	from huggingface_hub import snapshot_download
	from openunmix import predict, utils as umx_utils

	from pyharp import *


	model_card = ModelCard(
	name="Open-Unmix",
	description=(
	"Music source separation into vocals, drums, bass, and other "
	"instruments using deep neural networks. Repository: "
	"https://github.com/sigsep/open-unmix-pytorch Paper: "
	"https://doi.org/10.5281/zenodo.5069601 The code and the umxhq model "
	"are MIT, but the default umxl model carries CC-BY-NC-SA-4.0, so its "
	"results cannot be used commercially and must be shared alike."
	),
	author="F.-R. Stöter, S. Uhlich, A. Liutkus, Y. Mitsufuji",
	tags=[Subcategory.STEM_SEPARATION, SampleRate(44100), Channels(2), 'stems', 'domain:music'],
	)

	MODELS = ["umxl", "umxhq"]
	DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

	# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the stems
	# lossily again would expose the artifacts of the first encode and add a second set.
	LOSSLESS = {".wav", ".flac"}

	# Both separators are built and placed in the worker, which runs process_fn and starts with the
	# server, so that no GPU call builds one. The Wiener filter settings are plain attributes of a
	# separator, set for each request. openunmix downloads each checkpoint from Zenodo on first use,
	# which is too slow for a Space to start, so the same files are read from a mirror on the Hub,
	# linked where torch.hub looks before downloading.
	SEPARATORS = {}

	if __name__ != "__main__":
	_weights = Path(snapshot_download(
	"teamup-tech/Open-Unmix-weights", revision="2dca5d0f8b5692be0fa99c0d45c73fa7bff6d665"
	))
	_hub = Path(torch.hub.get_dir()) / "checkpoints"
	_hub.mkdir(parents=True, exist_ok=True)
	for _checkpoint in _weights.glob("/.pth"):
	with contextlib.suppress(FileExistsError):
	(_hub / _checkpoint.name).symlink_to(_checkpoint)

	for _model_name in MODELS:
	_separator = umx_utils.load_separator(model_str_or_path=_model_name, device=DEVICE,
	pretrained=True)
	_separator.freeze()
	SEPARATORS[_model_name] = _separator.to(DEVICE)


	# The whole mixture is separated, 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
	# reservation would exceed that is refused rather than started. A shorter one is left to
	# ZeroGPU, which admits it when the caller's own allowance covers it.
	MAX_GPU_SECONDS = 900


	def gpu_seconds(input_audio, model_name, niter, wiener_win_len):
	# Measured at about 2 s for 20 s of audio, 4 s for 60 s and 9 s for 300 s, upload included,
	# at any Wiener setting, padded. The separators are built when the worker starts, not inside
	# the call.
	return 5 + 0.04 * sf.info(input_audio).duration


	def process_fn(input_audio, model_name, niter, wiener_win_len):
	# Checked before the GPU is reserved, so a request the model cannot take spends none of the
	# caller's ZeroGPU allowance
	if model_name not in MODELS:
	raise gr.Error(f"Unknown model: {model_name}")
	if not input_audio:
	raise gr.Error("Provide the mixture to separate.")
	try:
	duration = sf.info(input_audio).duration
	except Exception as exc:
	raise gr.Error(f"Could not read the audio file: {exc}") from exc
	if duration == 0:
	raise gr.Error("The audio file is empty.")
	if gpu_seconds(input_audio, model_name, niter, wiener_win_len) > MAX_GPU_SECONDS:
	raise gr.Error(
	f"The mixture is {duration:.0f} s long, more than one request can separate in the "
	f"{MAX_GPU_SECONDS // 60} minutes it may run. Trim it."
	)
	return separate(input_audio, model_name, niter, wiener_win_len)


	# Inference mode as well as freeze(), since ZeroGPU recreates the parameters on the GPU for each
	# call, and the copies would require gradients again
	@spaces.GPU(duration=gpu_seconds)
	@torch.inference_mode()
	def separate(input_audio, model_name, niter, wiener_win_len):
	data, rate = sf.read(input_audio, dtype="float32", always_2d=True)
	audio = torch.from_numpy(data.T)

	# Open-Unmix separates at 44.1 kHz, so the session's rate is kept here and
	# the stems are resampled back to it before saving
	original_sr = rate
	original_channels = audio.shape[0]

	# A mono input is duplicated to feed the model, which separates in stereo. Two identical
	# channels leave the Wiener filter's spatial covariance singular, which past a couple of
	# iterations turns every stem to NaN, so the copy carries noise at -100 dBFS, below the
	# quietest step of a 16-bit file. The stems are averaged back to mono below.
	if audio.shape[0] == 1:
	noise = torch.randn(audio.shape, generator=torch.Generator().manual_seed(0))
	audio = torch.cat([audio, audio + 1e-5 * noise])

	separator = SEPARATORS[model_name]
	separator.niter = int(niter)
	separator.wiener_win_len = int(wiener_win_len)

	estimates = predict.separate(
	audio=audio,
	rate=rate,
	model_str_or_path=model_name,
	separator=separator,
	device=DEVICE
	)

	input_ext = Path(input_audio).suffix.lower()
	output_dir = Path(tempfile.mkdtemp())
	stems = {}
	for target in ["vocals", "drums", "bass", "other"]:
	target_tensor = estimates[target][0].cpu()

	if int(separator.sample_rate) != int(original_sr):
	target_tensor = torchaudio.functional.resample(
	target_tensor, int(separator.sample_rate), int(original_sr)
	)

	# A mono input was duplicated to feed the model, which separates in
	# stereo; giving back two identical-sounding channels would claim an
	# image the input never had
	if original_channels == 1:
	target_tensor = target_tensor.mean(dim=0, keepdim=True)

	stems[target] = target_tensor

	# Wiener filtering can push a stem past full scale, which a 16-bit file would flatten and a
	# FLAC file cannot hold at all, so when any stem exceeds it, all are scaled down together and
	# still sum to the mixture
	peak = max(float(stem.abs().max()) for stem in stems.values())
	output_files = {}
	for target, stem in stems.items():
	if peak > 1.0:
	stem = stem / peak
	out_file = output_dir / f"{target}{input_ext if input_ext in LOSSLESS else '.wav'}"
	sf.write(str(out_file), stem.numpy().T, int(original_sr))
	output_files[target] = str(out_file)

	return (
	output_files["vocals"],
	output_files["drums"],
	output_files["bass"],
	output_files["other"]
	)


	if __name__ == "__main__":
	with gr.Blocks() as demo:
	input_components = [
	gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
	"The mixture to separate. Mono input is duplicated to stereo first."
	)),
	gr.Dropdown(choices=MODELS, value="umxl", label="Model", info="umxl separates best, but its results cannot be used commercially. umxhq is MIT licensed."),
	gr.Slider(minimum=0, maximum=5, step=1, value=1, label="Wiener Filter Iterations", info="Number of iterations for Wiener filtering. 0 means softmask."),
	gr.Slider(minimum=100, maximum=500, step=50, value=300, label="Wiener Window Length", info="Number of frames on which to apply filtering independently."),
	]
	output_components = [
	gr.Audio(type="filepath", label="Vocals").set_info("The isolated vocal."),
	gr.Audio(type="filepath", label="Drums").set_info("The isolated drums."),
	gr.Audio(type="filepath", label="Bass").set_info("The isolated bass."),
	gr.Audio(type="filepath", label="Other").set_info((
	"Everything the other three stems did not take."
	)),
	]
	build_endpoint(
	model_card=model_card,
	input_components=input_components,
	output_components=output_components,
	process_fn=process_fn,
	)

	demo.queue().launch(show_error=True, pwa=True)