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)