import spaces # imported first, as ZeroGPU requires import sys sys.stdout.reconfigure(line_buffering=True) sys.setrecursionlimit(int(1e6)) import functools import tempfile import numpy as np import scipy.signal import torch import soundfile as sf import librosa import gradio as gr from pymixconsole.parameter import Parameter from pymixconsole.parameter_list import ParameterList from pyharp import ModelCard, build_endpoint from pyharp.tags import Channels, SampleRate, Subcategory from automix.common_dataprocessing import load_wav, save_wav, generate_data from config import * from automix.common_supernet import SuperNet from automix.common_datatypes import DataType from automix.utils_data_normalization import ( get_comp_matching, get_eq_matching, lufs_normalize, get_panning_matching, amp_to_db, get_reverb_send, ) model_card = ModelCard( name="FxNorm-Automix", description=( "Automatic music mixing with deep learning and out-of-domain data. " "Accepts up to 4 stems (vocals, bass, drums, other), each up to 3.5 " "minutes. Vocals are required, and bass, drums, and other are optional. The stems " "are mixed at 44.1 kHz in stereo, and the mix comes back at the highest rate among them. " "Repository: https://github.com/sony/FxNorm-automix Paper: " "https://arxiv.org/abs/2208.11428" ), author="Marco A. Martínez-Ramírez, Wei-Hsiang Liao, Giorgio Fabbro, Stefan Uhlich, Chihiro Nagashima, Yuki Mitsufuji", tags=[Subcategory.AUTOMATIC_MIXING, SampleRate(SR), Channels(N_CHANNELS), 'music', 'domain:music'], ) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" USE_AMP = torch.cuda.is_available() EQ_PARAMETERS = ParameterList() EQ_PARAMETERS.add(Parameter('low_shelf_gain', EQ_GAIN, 'float', minimum=EQ_GAIN, maximum=EQ_GAIN)) EQ_PARAMETERS.add(Parameter('low_shelf_freq', 600.0, 'float', minimum=500.0, maximum=700.0)) EQ_PARAMETERS.add(Parameter('high_shelf_gain', EQ_GAIN, 'float', minimum=EQ_GAIN, maximum=EQ_GAIN)) EQ_PARAMETERS.add(Parameter('high_shelf_freq', 8500.0, 'float', minimum=7000.0, maximum=10000.0)) PRE_REVERB_PARAMETERS = ParameterList() PRE_REVERB_PARAMETERS.add(Parameter('wet', 0.5, 'float', minimum=0.5, maximum=0.5)) PRE_REVERB_PARAMETERS.add(Parameter('dry', 0.5, 'float', minimum=0.5, maximum=0.5)) PRE_REVERB_PARAMETERS.add(Parameter('decay', 1.0, 'float', minimum=1.0, maximum=1.0)) PRE_REVERB_PARAMETERS.add(Parameter('pre_delay', 0, 'int', units='ms', minimum=0, maximum=0)) REVERB_PARAMETERS = ParameterList() REVERB_PARAMETERS.add(Parameter('wet', 1.0, 'float', minimum=1.0, maximum=1.0)) REVERB_PARAMETERS.add(Parameter('dry', 0.0, 'float', minimum=0.0, maximum=0.0)) REVERB_PARAMETERS.add(Parameter('decay', 1.0, 'float', minimum=1.0, maximum=1.0)) REVERB_PARAMETERS.add(Parameter('pre_delay', 0, 'int', units='ms', minimum=0, maximum=0)) def _load_ir(path): data = load_wav(path) return {'impulse_response': functools.partial(generate_data, file_path_or_data=data)} # IR and pre-IR are used from OpenAIR library # Abies Grandis Forest, Wheldrake Wood (s1r1 and s1r2) # Lady Chapel, St Albans Cathedral (stalbans_a_ortf, stalbans_b_ortf) # Stairway, University of York (stairwell_ortf) PRE_IR = [_load_ir('pre-ir/S1R1_SF.wav'), _load_ir('pre-ir/S1R2_SF.wav')] IR = [_load_ir('ir/stalbans_a_ortf.wav'), _load_ir('ir/stalbans_b_ortf.wav'), _load_ir('ir/stairwell_ortf.wav')] PRE_REVERB_PARAMETERS.add(Parameter('index', 0, 'int', minimum=0, maximum=len(PRE_IR))) REVERB_PARAMETERS.add(Parameter('index', 0, 'int', minimum=0, maximum=len(IR))) # --- Load model --- # pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the # server, so the model is built there alone and placed on DEVICE on its main thread, where # ZeroGPU's emulation records the placement, rather than inside a GPU call if __name__ != "__main__": net = torch.load('checkpoints/net_mixture.dump', map_location='cpu', weights_only=False) net.load_state_dict(torch.load('checkpoints/current_model_for_mixture.params', map_location='cpu', weights_only=False)) super_net = SuperNet( net, stft_window=torch.from_numpy(STFT_WINDOW.astype(np.float32)), stft_hop_length=HOP_LENGTH_CONFIG, batched_valid=False, unfolding_params=None, training_length=TRAINING_SEQ_LENGTH, training_batch_size=BATCH_SIZE, use_amp=USE_AMP, ) super_net.eval() super_net.to(DEVICE) # --- Load and smooth features --- features_mean = np.load('checkpoints/features_MUSDB18.npy', allow_pickle=True)[()] for _effect in EFFECTS: if _effect == 'eq': for _key in STEMS: window_length = 401 if _key in ['other', 'vocals'] else 151 features_mean[_effect][_key] = scipy.signal.savgol_filter( features_mean[_effect][_key], window_length, 1, mode='mirror') elif _effect == 'panning': for _key in STEMS: features_mean[_effect][_key] = scipy.signal.savgol_filter( features_mean[_effect][_key], 501, 1, mode='mirror') def normalize_audio_wave(stem_args): """ Effect Normalizes audio input Args: tuple: audio x list of effects to normalize x source (e.g., vocals) x dataset features Returns: normalized audio """ audio = stem_args[0] effects = stem_args[1] src = stem_args[2] features_mean_ = stem_args[3] print(f'{src} normalizing {effects}...') all_zeros = not np.any(audio) if not all_zeros: audio_track = np.pad(audio, ((FFT_SIZE, FFT_SIZE), (0, 0))) assert len(audio_track.shape) == 2 if audio_track.shape[1] == 1: audio_track = np.repeat(audio_track, 2, axis=-1) output_audio = audio_track.copy() max_db = amp_to_db(np.max(np.abs(output_audio))) if max_db > MIN_DB: for effect in effects: if effect == 'eq': for ch in range(audio_track.shape[1]): audio_eq_matched = get_eq_matching( output_audio[:, ch], features_mean_[effect][src], sr=SR, n_fft=FFT_SIZE, hop_length=HOP_LENGTH, min_db=MIN_DB, ntaps=NTAPS, lufs=LUFS) np.copyto(output_audio[:, ch], audio_eq_matched) elif effect == 'compression': assert len(features_mean_[effect][src]) == 2 for ch in range(audio_track.shape[1]): audio_comp_matched = get_comp_matching( output_audio[:, ch], features_mean_[effect][src][0], features_mean_[effect][src][1], comp_settings[src]['ratio'], comp_settings[src]['attack'], comp_settings[src]['release'], sr=SR, min_db=MIN_DB, min_th=COMP_MIN_TH, comp_peak_norm=COMP_PEAK_NORM, max_ratio=COMP_MAX_RATIO, n_mels=comp_settings[src]['n_mels'], true_peak=COMP_TRUE_PEAK, percentile=COMP_PERCENTILE, expander=COMP_USE_EXPANDER) np.copyto(output_audio[:, ch], audio_comp_matched[:, 0]) elif effect == 'panning': if PANNING_FRAMES: ref_phi = features_mean_[effect][src][:-1] bins_per_frame = int(2 * (features_mean_[effect][src].shape[0] - 1) / FFT_SIZE_PAN) ref_phi = ref_phi.reshape(-1, bins_per_frame).mean(axis=1) ref_phi = np.append(ref_phi, features_mean_[effect][src][-1]) else: ref_phi = features_mean_[effect][src] MAX_FREQ_PAN = 6000 if src == 'drums' else 16000 audio_pan_matched = get_panning_matching(output_audio, PANNING_GAIN_K * ref_phi, sr=SR, n_fft=FFT_SIZE_PAN, hop_length=HOP_LENGTH_PAN, min_db_f=MIN_DB_f, max_freq_pan=MAX_FREQ_PAN, frames=PANNING_FRAMES) np.copyto(output_audio, audio_pan_matched) elif effect == 'loudness': output_audio = lufs_normalize(output_audio, SR, features_mean_[effect][src]) elif effect == 'reverb': if src not in ['bass', 'drums']: audio_reverb_send = get_reverb_send( output_audio, EQ_PARAMETERS, REVERB_PARAMETERS, IR, bands=EQ_BANDS_VOCALS_OTHER, eq_prob=EQ_PROBABILITY, rv_prob=REVERB_PROBABILITY, parallel=FX_PARALLEL, shuffle=FX_SHUFFLE, sr=SR) np.copyto(output_audio, audio_reverb_send) elif effect == 'prereverb': if src not in ['bass', 'drums']: audio_reverb_send = get_reverb_send( output_audio, EQ_PARAMETERS, PRE_REVERB_PARAMETERS, PRE_IR, bands=EQ_BANDS_VOCALS_OTHER, eq_prob=PRE_EQ_PROBABILITY, rv_prob=PRE_REVERB_PROBABILITY, parallel=PRE_FX_PARALLEL, shuffle=PRE_FX_SHUFFLE, sr=SR) np.copyto(output_audio, audio_reverb_send) output_audio = output_audio[FFT_SIZE:FFT_SIZE + audio.shape[0]] else: print(f'{src} is only zeros...') output_audio = audio return output_audio def process_fn(vocals_path, bass_path, drums_path, other_path): # Checked before the GPU is reserved, so stems the model cannot take spend none of the # caller's ZeroGPU allowance paths = {'vocals': vocals_path, 'bass': bass_path, 'drums': drums_path, 'other': other_path} if not any(paths.values()): raise gr.Error("Please provide at least one stem.") for stem in STEMS: if paths[stem]: try: info = sf.info(paths[stem]) except Exception as exc: raise gr.Error(f"{stem}: could not read the audio file: {exc}") from exc if info.frames == 0: raise gr.Error(f"{stem}: the audio file is empty.") if info.channels not in (1, N_CHANNELS): raise gr.Error(f'{stem}: expected {N_CHANNELS} channels, got {info.channels}') if info.duration > MAX_STEM_SECONDS: raise gr.Error( f"{stem}: the stem is {info.duration:.1f} s long. Trim it to " f"{MAX_STEM_SECONDS / 60:g} minutes or less." ) return mix(paths) # Every stem is padded with 30 s of context, and the model takes at most MAX_LENGTH seconds with # the padding, so a longer stem would come back cut short CONTEXT_SECONDS = 30 MAX_STEM_SECONDS = MAX_LENGTH - CONTEXT_SECONDS def gpu_seconds(paths): # Measured at about 29 s for four 30 s stems and for one 60 s stem, padded longest = max(sf.info(path).duration for path in paths.values() if path) return 15 + 0.4 * (longest + CONTEXT_SECONDS) @spaces.GPU(duration=gpu_seconds) def mix(paths): # --- Load and validate stems --- pad_samples = CONTEXT_SECONDS * SR audio = {} samples = [] subtype_list = [] samplingrate_list = [] for stem in STEMS: if paths[stem]: # soundfile reads every format HARP sends, where the authors' reader takes only # integer PCM WAV, and returns the integer PCM the pipeline below expects _arr, _sr = sf.read(paths[stem], dtype="int32", always_2d=True) samplingrate_list.append(_sr) # A lossy or floating-point input has no PCM depth for the WAV output to keep subtype = sf.info(paths[stem]).subtype subtype_list.append(subtype if subtype.startswith("PCM_") else "PCM_24") # ensure stereo if _arr.shape[1] == 1: _arr = np.repeat(_arr, 2, axis=-1) elif _arr.shape[1] != N_CHANNELS: raise gr.Error(f'{stem}: expected {N_CHANNELS} channels, got {_arr.shape[1]}') if _sr != SR: # resample to model's expected SR (44100); librosa needs float32 in [-1, 1] # so we convert int16 to float, resample, then convert back to int16 temp = _arr.astype(np.float32) / (1. + np.iinfo(_arr.dtype).max) temp = librosa.resample(temp.T, orig_sr=_sr, target_sr=SR, res_type='kaiser_best').T temp = temp * (1 + np.iinfo(np.int16).max) _arr = np.clip(temp, np.iinfo(np.int16).min, np.iinfo(np.int16).max).astype(np.int16) # Pad the start with 30s of wrapped audio so the model has context at # t=0. This follows the resample rather than preceding it, because # pad_samples counts samples at SR: padding first adds that many # samples at the input's own rate, and resampling then leaves a pad # shorter than the trim after inference removes. A 48 kHz input came # back empty for exactly that reason. _arr = np.pad(_arr, [(pad_samples, 0), (0, 0)], 'wrap') samples.append(_arr.shape[0]) audio[stem] = (SR, _arr) else: # missing stems are silent so the model always receives 4 tracks audio[stem] = (SR, np.zeros((SR, N_CHANNELS), dtype=np.int16)) max_samples = min(max(samples), int(MAX_LENGTH * SR)) max_samplingrate = max(samplingrate_list) max_subtype = max(subtype_list) # --- Build input tensor --- data = np.zeros((len(STEMS), 1, max_samples, N_CHANNELS), dtype=np.float32) for k, stem in enumerate(STEMS): stem_audio = generate_data(audio[stem]) if stem_audio[:max_samples].shape[0] < max_samples: zeros_ = np.zeros_like(data[k][0]) zeros_[:stem_audio[:max_samples].shape[0]] = stem_audio[:max_samples] data[k] = zeros_[:max_samples] else: data[k] = stem_audio[:max_samples] # round up to the next multiple of KERNEL_SIZE_ENCODER as required by the encoder new_samples = (1 + max_samples // KERNEL_SIZE_ENCODER) * KERNEL_SIZE_ENCODER data = np.pad(data, [(0, 0), (0, 0), (0, new_samples - max_samples), (0, 0)]) # --- Effect normalization --- # changed from org. inference #512 because parallelization may not be safe on HF's CPU tier for k, stem in enumerate(STEMS): data[k][0] = normalize_audio_wave((data[k][0], EFFECTS, stem, features_mean)) # --- Run model inference --- test_data = torch.from_numpy(data).to(DEVICE) with torch.no_grad(): test_out = super_net.inference(test_data) audio_out = test_out[DataType.TIME_SAMPLES].cpu().numpy() audio_out = audio_out[..., :max_samples] # trim encoder padding audio_out = audio_out[..., pad_samples:] # trim the 30s context pad added at load # first output stem is the mix audio_output = audio_out[0, 0, ...] # --- Resample and save output --- if max_samplingrate != SR: audio_output = librosa.resample(audio_output, orig_sr=SR, target_sr=max_samplingrate, res_type='kaiser_best') with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f: out_path = f.name save_wav(out_path, max_samplingrate, audio_output.T, subtype=max_subtype) return out_path if __name__ == "__main__": with gr.Blocks(title="FxNorm Automix") as demo: input_components = [ gr.Audio(type="filepath", label="Vocals").harp_required(True).set_info("The vocal stem, up to 3.5 minutes. Required."), gr.Audio(type="filepath", label="Bass").harp_required(False).set_info(( "The bass stem. Optional, and left out if the song has none." )), gr.Audio(type="filepath", label="Drums").harp_required(False).set_info("The drum stem. Optional."), gr.Audio(type="filepath", label="Other").harp_required(False).set_info(( "Everything else: guitars, keys, strings. Optional." )), ] output_components = [ gr.Audio(type="filepath", label="Mix").set_info("Automatic mix of the input stems."), ] 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)